diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 94213b82efb5..c21fc9ceb6af 100644 --- a/docs/source/features/speculative-decoding.md +++ b/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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/layers.py index 0652ede91bb8..dd23d60e5304 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/layers.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 6346762c72f3..5defbd078e2d 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/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: diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index e9a8683c9282..d5a7346811a2 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/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: diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 3792c621690f..3e7f42359151 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/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( diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 744a82a1d361..d7322ce9f3d2 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/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]( diff --git a/tensorrt_llm/_torch/speculative/fused_accept.py b/tensorrt_llm/_torch/speculative/fused_accept.py index dc0609700ae7..aaa93721d5e2 100644 --- a/tensorrt_llm/_torch/speculative/fused_accept.py +++ b/tensorrt_llm/_torch/speculative/fused_accept.py @@ -165,6 +165,41 @@ def fused_strict_accept( return accepted, num_accepted +@triton.jit +def _step0_gather_id( + seq_lens_ptr, num_accepted_ptr, i, num_contexts, num_ctx_tokens, K, BLOCK_C: tl.constexpr +): + """Step 0's row of request ``i`` and its accepted-token count.""" + is_ctx = i < num_contexts + offs = tl.arange(0, BLOCK_C) + lens = tl.load(seq_lens_ptr + offs, mask=(offs <= i) & is_ctx, other=0).to(tl.int64) + acc = tl.load(num_accepted_ptr + i).to(tl.int64) + # contexts: cumsum(seq_lens)[i] - 1 (last prompt token); generations: last accepted token of the block + gid = tl.where( + is_ctx, + tl.sum(lens, axis=0) - 1, + num_ctx_tokens + (i - num_contexts).to(tl.int64) * (K + 1) + acc - 1, + ) + return gid, acc + + +@triton.jit +def _step0_gather_ids_kernel( + seq_lens_ptr, + num_accepted_ptr, + ids_out_ptr, + num_contexts, + num_ctx_tokens, + K: tl.constexpr, + BLOCK_C: tl.constexpr, +): + i = tl.program_id(0) + gid, _ = _step0_gather_id( + seq_lens_ptr, num_accepted_ptr, i, num_contexts, num_ctx_tokens, K, BLOCK_C + ) + tl.store(ids_out_ptr + i, gid) + + @triton.jit def _step0_epilogue_kernel( seq_lens_ptr, @@ -190,25 +225,25 @@ def _step0_epilogue_kernel( BLOCK_C: tl.constexpr, BLOCK_H: tl.constexpr, LAST: tl.constexpr, + # The draft layer kept only the gather-id rows: ``ids_out_ptr`` already holds the ids (read, not + # written) and the hidden rows need no gather (``hid_ptr`` / ``hid_out_ptr`` unused). + PICKED: tl.constexpr, ): """One program per request ``i`` (every write is to row ``i`` only; ``seq_lens_ptr`` is never written).""" i = tl.program_id(0) - is_ctx = i < num_contexts - offs = tl.arange(0, BLOCK_C) - lens = tl.load(seq_lens_ptr + offs, mask=(offs <= i) & is_ctx, other=0).to(tl.int64) - acc = tl.load(num_accepted_ptr + i).to(tl.int64) - # contexts: cumsum(seq_lens)[i] - 1 (last prompt token); generations: last accepted token of the block - gid = tl.where( - is_ctx, - tl.sum(lens, axis=0) - 1, - num_ctx_tokens + (i - num_contexts).to(tl.int64) * (K + 1) + acc - 1, - ) - dkv = tl.where(is_ctx, 1, acc - K) - tl.store(ids_out_ptr + i, gid) - h = tl.arange(0, BLOCK_H) - hm = h < H - v = tl.load(hid_ptr + gid * hid_stride + h, mask=hm) - tl.store(hid_out_ptr + i.to(tl.int64) * H + h, v, mask=hm) + if PICKED: + gid = tl.load(ids_out_ptr + i) + acc = tl.load(num_accepted_ptr + i).to(tl.int64) + else: + gid, acc = _step0_gather_id( + seq_lens_ptr, num_accepted_ptr, i, num_contexts, num_ctx_tokens, K, BLOCK_C + ) + tl.store(ids_out_ptr + i, gid) + h = tl.arange(0, BLOCK_H) + hm = h < H + v = tl.load(hid_ptr + gid * hid_stride + h, mask=hm) + tl.store(hid_out_ptr + i.to(tl.int64) * H + h, v, mask=hm) + dkv = tl.where(i < num_contexts, 1, acc - K) for r in tl.static_range(R): p = tl.load(pos_ptr + r * pos_row_stride + gid) for j in tl.static_range(1, K): @@ -254,11 +289,15 @@ def step0_epilogue( num_ctx_tokens: int, accepted_tokens: Optional[torch.Tensor] = None, last_out: Optional[torch.Tensor] = None, + picked_gather_ids: Optional[torch.Tensor] = None, ) -> Optional[Step0Epilogue]: """``seq_lens`` must not alias ``seq_lens_out`` (read the pre-draft-loop copy). With ``accepted_tokens`` ``[b, >= draft_len + 1]`` and ``last_out`` (contiguous ``[b]``), also ``last_out[i] = - accepted_tokens[i, num_accepted[i] - 1]``. Returns None when the shapes are not supported (caller keeps the - stock path).""" + accepted_tokens[i, num_accepted[i] - 1]``. ``picked_gather_ids`` (the ``step0_gather_ids`` the draft + layer was given as ``pick_rows``): ``hidden_states`` holds only those rows (``[b, H]``, in request order), + and the result reuses both instead of gathering. Returns None when the shapes are not supported (caller + keeps the stock path).""" + picked = picked_gather_ids is not None if ( seq_lens.data_ptr() == seq_lens_out.data_ptr() or hidden_states.dim() != 2 @@ -269,14 +308,20 @@ def step0_epilogue( or draft_len < 1 ): return None + if picked: + assert hidden_states.shape[0] == batch_size and picked_gather_ids.shape == (batch_size,) + assert picked_gather_ids.dtype == torch.int64 and picked_gather_ids.is_contiguous() pos2 = position_ids.reshape(-1, position_ids.shape[-1]) if pos2.dim() != 2 or (pos2.shape[0] > 1 and pos2.stride(0) < pos2.shape[1]): return None rows = pos2.shape[0] h = hidden_states.shape[1] dev = hidden_states.device - ids = torch.empty((batch_size,), dtype=torch.int64, device=dev) - hid = torch.empty((batch_size, h), dtype=hidden_states.dtype, device=dev) + if picked: + ids, hid = picked_gather_ids, hidden_states + else: + ids = torch.empty((batch_size,), dtype=torch.int64, device=dev) + hid = torch.empty((batch_size, h), dtype=hidden_states.dtype, device=dev) pos_all = torch.empty( (max(draft_len - 1, 1), rows, batch_size), dtype=position_ids.dtype, device=dev ) @@ -311,14 +356,42 @@ def step0_epilogue( BLOCK_C=max(16, triton.next_power_of_2(max(num_contexts, 1))), BLOCK_H=triton.next_power_of_2(h), LAST=last, + PICKED=picked, num_warps=4, ) return Step0Epilogue(ids, hid, pos_all, tuple(position_ids.shape[:-1])) +def step0_gather_ids( + seq_lens: torch.Tensor, + num_accepted: torch.Tensor, + num_contexts: int, + batch_size: int, + draft_len: int, + num_ctx_tokens: int, +) -> Optional[torch.Tensor]: + """``step0_epilogue``'s ``gather_ids`` alone (``[b]`` int64), for use before the step-0 draft forward. + Returns None when the batch is not supported (caller computes them with the stock ops).""" + if batch_size < 1 or batch_size > 1024 or draft_len < 1: + return None + ids = torch.empty((batch_size,), dtype=torch.int64, device=num_accepted.device) + _step0_gather_ids_kernel[(batch_size,)]( + seq_lens, + num_accepted, + ids, + num_contexts, + num_ctx_tokens, + K=draft_len, + BLOCK_C=max(16, triton.next_power_of_2(max(num_contexts, 1))), + num_warps=1, + ) + return ids + + __all__ = [ "MAX_DRAFT_LEN", "fused_strict_accept", "step0_epilogue", + "step0_gather_ids", "Step0Epilogue", ] diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d68ebfff9b8d..f6745310bdcf 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -14,6 +14,7 @@ # limitations under the License. import copy +import inspect import os from abc import ABC, abstractmethod from contextlib import contextmanager @@ -338,6 +339,68 @@ def get_force_num_accepted_tokens_float() -> float: return 0.0 +DISABLE_DRAFT_ROW_PICK_ENV_VAR = "TLLM_SPEC_DISABLE_DRAFT_ROW_PICK" + + +def draft_layer_picks_rows(layer: Optional[nn.Module], mapping=None) -> bool: + """Whether ``layer`` is fed the rows to keep after its attention. + + A draft layer opts in with ``supports_pick_rows = True`` and a + ``pick_rows`` forward argument (``[n]`` indices into its input rows). + Given ``pick_rows``, the layer's attention still runs on every input row + (writing the draft KV of each of them), and the rest of the layer (the + residual, norms, MLP / MoE and final norm, all per-token) runs on those + ``n`` rows only, so it returns ``n`` rows: the input rows ``pick_rows`` + would have selected from the full output. Under attention DP the caller + passes the per-rank picked row counts as ``all_rank_num_tokens``, which a + supporting layer uses only after the pick. + + Declined: context parallelism (sharded attention outputs), layers without + the attribute, and layers whose own ``forward`` does not name + ``pick_rows``. The last guards subclasses of a supporting layer that + override ``forward`` and would otherwise inherit the attribute; a + subclass that keeps its own row selection should also set + ``supports_pick_rows = False``. ``TLLM_SPEC_DISABLE_DRAFT_ROW_PICK=1`` + turns picking off. + """ + if layer is None or not getattr(layer, "supports_pick_rows", False): + return False + if "pick_rows" not in inspect.signature(type(layer).forward).parameters: + return False + if mapping is not None and mapping.cp_size > 1: + return False + return os.environ.get(DISABLE_DRAFT_ROW_PICK_ENV_VAR, "0") != "1" + + +class DraftRowPick: + """A worker's ``draft_layer_picks_rows`` decision for its draft layers, + resolved again (and logged) whenever their classes change: a model may + re-class its draft layers after the worker is built.""" + + def __init__(self, layers_name: str): + self._layers_name = layers_name + self._layer_classes = None + self._picks_rows = False + + def resolve(self, + layers: List[Optional[nn.Module]], + mapping=None, + decline: Optional[str] = None) -> bool: + """``decline``: the reason this worker cannot pick (None: it can).""" + layer_classes = tuple(type(layer) for layer in layers) + if layer_classes != self._layer_classes: + self._layer_classes = layer_classes + self._picks_rows = (decline is None and len(layers) > 0 and all( + draft_layer_picks_rows(layer, mapping) for layer in layers)) + names = ", ".join(sorted({c.__name__ + for c in layer_classes})) or "none" + logger.info( + f"{self._layers_name} ({names}) run after attention on " + + ("one row per request" if self._picks_rows else "every row") + + (f" ({decline})" if decline is not None else "")) + return self._picks_rows + + class SpeculativeDecodingMode(IntEnum): # Values are explicit and must stay stable: this enum is in the # cross-process serialization allowlist (see llmapi/serialization.py), so a diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 50de723c3968..6ef44b5e0e6c 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -7,13 +7,14 @@ import torch from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger from ..attention.backends import AttentionMetadata from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.scheduler import ScheduledRequests -from .interface import SpecMetadata, SpecWorkerBase +from .interface import DraftRowPick, SpecMetadata, SpecWorkerBase from .sa_enhancer import SADraftEnhancer if TYPE_CHECKING: @@ -271,6 +272,7 @@ def __init__(self, model_config, "mapping", None) self.is_thop = False self._is_mamba_hybrid_cache = None + self._row_pick = DraftRowPick("MTP layers") self.sa_enhancer: Optional[SADraftEnhancer] = None if spec_config.sa_config is not None: self.sa_enhancer = SADraftEnhancer(spec_config.sa_config.threshold) @@ -279,6 +281,16 @@ def __init__(self, def max_draft_len(self) -> int: return self.spec_config.max_draft_len + def _mtp_layers_pick_rows(self, draft_model) -> bool: + """Whether the MTP layers take ``pick_rows`` (see + draft_layer_picks_rows). Declined under attention DP, where this + worker passes the MTP layers no per-rank token counts.""" + adp = self.mapping is not None and self.mapping.enable_attention_dp + return self._row_pick.resolve(list( + getattr(draft_model, "mtp_layers", None) or ()), + self.mapping, + decline="attention DP" if adp else None) + def _commit_target_mamba_states( self, attn_metadata: AttentionMetadata, @@ -473,6 +485,16 @@ def _forward_impl( draft_kv_cache_manager = self.get_draft_kv_cache_manager( resource_manager) + # Each MTP module reads its output on the last row of every request + # only (the LM-head input and the next module's carried hidden state), + # so a supporting layer runs its post-attention part on those rows. + picks_rows = self._mtp_layers_pick_rows(draft_model) + if picks_rows: + logger.info_once( + f"MTP layers pick {batch_size} of " + f"{draft_inputs['input_ids'].shape[0]} rows after attention", + key="spec_mtp_pick_rows") + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i, mtp_layer in enumerate( draft_model.mtp_layers[:runtime_draft_len]): @@ -482,12 +504,23 @@ def _forward_impl( num_accepted_tokens, draft_step=i) + if picks_rows: + draft_inputs["pick_rows"] = last_tokens_idx hidden_states = mtp_layer(embed_tokens=draft_model.embed_tokens, **draft_inputs) - - logits = mtp_layer.shared_head(hidden_states, - draft_model.lm_head, - attn_metadata).float() + # The rows of the requests' last tokens, in request order. + last_rows = (slice(None) if picks_rows else last_tokens_idx) + + if picks_rows: + # The output holds only the last rows: the head takes + # every row it is given (return_context_logits). + logits = mtp_layer.shared_head(hidden_states, + draft_model.lm_head, + attn_metadata, True).float() + else: + logits = mtp_layer.shared_head(hidden_states, + draft_model.lm_head, + attn_metadata).float() if self.guided_decoder is not None: self.guided_decoder.execute_draft_batch(logits, draft_step=i) @@ -504,7 +537,7 @@ def _forward_impl( draft_hidden_states = draft_inputs["hidden_states"] draft_hidden_states[:-1] = draft_hidden_states[1:].clone() draft_hidden_states[last_tokens_idx] = hidden_states[ - last_tokens_idx, :] + last_rows, :] draft_inputs = { "input_ids": input_ids, "position_ids": draft_inputs["position_ids"], diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b55fc2dc677b..d0828f799fd7 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -340,6 +340,8 @@ l0_b200: - unittest/_torch/kernels/test_spec_step_copies.py - unittest/_torch/pyexecutor/engine/runners/decoder/test_inputs.py -k "test_spec_decode_graph_step_gather_kernel" - unittest/_torch/speculative/test_fused_accept.py + # Draft step 0 runs the draft layer after attention only on the rows it drafts from. + - unittest/_torch/speculative/test_draft_row_pick.py - unittest/_torch/pyexecutor/test_steady_spec_prepare.py - unittest/_torch/modules/test_mxfp8_draft_head.py # Shared fused_qk_norm_rope cells used by the K3 drafter. diff --git a/tests/unittest/_torch/speculative/test_draft_row_pick.py b/tests/unittest/_torch/speculative/test_draft_row_pick.py new file mode 100644 index 000000000000..a25090358c9b --- /dev/null +++ b/tests/unittest/_torch/speculative/test_draft_row_pick.py @@ -0,0 +1,595 @@ +# 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. +"""Draft step 0 runs the draft layer's post-attention part only on the rows +the draft loop reads (``pick_rows``). + +The draft tokens and the hidden states carried to step 1 must match the +unpicked run, the attention (draft KV) must still see every row, and a draft +layer without ``pick_rows`` support must run exactly as before. +""" + +import contextlib +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import DraftRowPick, draft_layer_picks_rows +from tensorrt_llm._torch.speculative.mtp import MTPWorker + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +H = 64 +VOCAB = 97 + + +class _Meta: + """The attention-metadata surface the linear draft loop touches.""" + + def __init__(self, seq_lens, num_contexts, kv_lens): + self._seq_lens = torch.tensor(seq_lens, dtype=torch.int32) + self._seq_lens_cuda = self._seq_lens.cuda() + self.kv_lens_cuda = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + self.num_contexts = num_contexts + self.num_seqs = len(seq_lens) + self.num_ctx_tokens = int(sum(seq_lens[:num_contexts])) + self.kv_cache_manager = None + self.enable_flash_mla = False + self.use_spec_decoding = True + self.all_rank_num_tokens = None + + @property + def seq_lens_cuda(self): + return self._seq_lens_cuda + + def on_update(self): + pass + + def update_for_spec_dec(self): + pass + + +class _DraftAttnLayer(nn.Module): + """A toy draft layer: "attention" mixes every input row (a running mean + over the flattened rows) and is logged as the draft KV it writes; the rest + is per-token (residual, MLP, norm).""" + + def __init__(self, seed): + super().__init__() + g = torch.Generator().manual_seed(seed) + self.embed = (torch.randn(VOCAB, H, generator=g) * 0.5).cuda() + self.w = (torch.randn(H, H, generator=g) / H**0.5).cuda() + self.head = torch.randn(H, VOCAB, generator=g).cuda() + self.kv_writes = [] # every row the attention saw, per call + self.post_rows = [] # rows the post-attention part ran on, per call + self.inputs = [] # hidden input rows, per call + self.all_rank_num_tokens = [] # the all_rank_num_tokens argument, per call + + def _attention(self, input_ids, position_ids, hidden_states): + x = ( + hidden_states + + self.embed[input_ids.long()] + + 0.01 * (position_ids.reshape(-1, 1).float()) + ) + self.kv_writes.append(x.clone()) + n = torch.arange(1, x.shape[0] + 1, device=x.device).unsqueeze(1) + return x.cumsum(0) / n, x + + def _post(self, attn_out, residual): + self.post_rows.append(attn_out.shape[0]) + y = torch.tanh(attn_out @ self.w) + residual + return y / y.norm(dim=-1, keepdim=True) + + def shared_head(self, hidden_states, lm_head, attn_metadata, return_context_logits=False): + if not return_context_logits: # the request's last row, as DeepseekV3MTPHead + last = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + hidden_states = hidden_states[last] + return hidden_states @ self.head + + +class _PickingMtpLayer(_DraftAttnLayer): + supports_pick_rows = True + + def forward( + self, + input_ids, + position_ids, + hidden_states, + embed_tokens, + attn_metadata, + all_rank_num_tokens=None, + spec_metadata=None, + pick_rows=None, + ): + self.inputs.append(hidden_states.clone()) + self.all_rank_num_tokens.append(all_rank_num_tokens) + attn_out, residual = self._attention(input_ids, position_ids, hidden_states) + if pick_rows is not None: + attn_out, residual = attn_out[pick_rows], residual[pick_rows] + return self._post(attn_out, residual) + + +class _PlainMtpLayer(_DraftAttnLayer): + """No ``pick_rows`` support: must never be given one.""" + + def forward( + self, + input_ids, + position_ids, + hidden_states, + embed_tokens, + attn_metadata, + all_rank_num_tokens=None, + spec_metadata=None, + **kwargs, + ): + assert "pick_rows" not in kwargs + self.inputs.append(hidden_states.clone()) + self.all_rank_num_tokens.append(all_rank_num_tokens) + attn_out, residual = self._attention(input_ids, position_ids, hidden_states) + return self._post(attn_out, residual) + + +class _OverridingMtpLayer(_PickingMtpLayer): + """Inherits the capability attribute but overrides ``forward`` without + ``pick_rows``: must be treated as unsupported.""" + + def forward( + self, input_ids, position_ids, hidden_states, embed_tokens, attn_metadata, **kwargs + ): + assert "pick_rows" not in kwargs + return super().forward(input_ids, position_ids, hidden_states, embed_tokens, attn_metadata) + + +class _Eagle3Model(_DraftAttnLayer): + """Eagle3DraftModel-style: returns (hidden, hidden_to_save).""" + + def __init__(self, seed, supports): + super().__init__(seed) + self.supports_pick_rows = supports + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + hidden_states, + spec_metadata, + all_rank_num_tokens=None, + pick_rows=None, + ): + assert pick_rows is None or self.supports_pick_rows + self.inputs.append(hidden_states.clone()) + self.all_rank_num_tokens.append(all_rank_num_tokens) + attn_out, residual = self._attention(input_ids, position_ids, hidden_states) + if pick_rows is not None: + attn_out, residual = attn_out[pick_rows], residual[pick_rows] + y = self._post(attn_out, residual) + return y, y + residual + + +def _worker(is_mtp_eagle, fast_paths): + w = object.__new__(Eagle3OneModelWorker) + w.is_mtp_eagle = is_mtp_eagle + w.spec_decode_fast_paths = fast_paths + w.guided_decoder = None + w.sa_enhancer = None + w.model_config = None + w.mapping = None + w.draft_head = None + w._d2t = None + w._draft_folds_kv_lens_increment = False + w._step0_row_pick = DraftRowPick("test draft layer") + w.draft_kv_cache_context = lambda md, mgr: contextlib.nullcontext(md) + w.sample_draft_tokens = ( + lambda logits, + spec_metadata, + batch_size, + draft_step, + mapping_lm_head_tp=None: logits.argmax(-1).to(torch.int32) + ) + return w + + +def _run_loop( + layer, is_mtp_eagle, fast_paths, num_contexts, num_gens, k, seed, all_rank=None, subseq=None +): + g = torch.Generator().manual_seed(seed) + ctx_lens = torch.randint(2, 9, (num_contexts,), generator=g).tolist() + seq_lens = ctx_lens + [k + 1] * num_gens + b = num_contexts + num_gens + n = sum(seq_lens) + num_acc = torch.cat( + [ + torch.ones(num_contexts, dtype=torch.int32), + torch.randint(1, k + 2, (num_gens,), generator=g, dtype=torch.int32), + ] + ).cuda() + accepted = torch.randint(0, VOCAB, (b, k + 1), generator=g, dtype=torch.int32).cuda() + md = _Meta(seq_lens, num_contexts, torch.randint(50, 90, (b,), generator=g).tolist()) + inputs = { + "input_ids": torch.randint(0, VOCAB, (n,), generator=g, dtype=torch.int32).cuda(), + "position_ids": torch.randint(0, 1000, (n,), generator=g, dtype=torch.int32).cuda(), + "hidden_states": torch.randn(n, H, generator=g).cuda(), + "attn_metadata": md, + "spec_metadata": None, + } + spec_metadata = SimpleNamespace( + runtime_draft_len=k, + batch_indices_cuda=torch.arange(b, dtype=torch.int32, device="cuda"), + wants_advanced_draft_sampling=False, + all_rank_num_tokens=all_rank, + subseq_all_rank_num_tokens=subseq, + max_num_requests=b, + ) + inputs["spec_metadata"] = spec_metadata + if is_mtp_eagle: + draft_model = SimpleNamespace(mtp_layers=[layer], embed_tokens=None, lm_head=None) + else: + draft_model = SimpleNamespace( + model=layer, + lm_head=None, + logits_processor=lambda h, lm_head, md, ctx: layer.shared_head(h, lm_head, md, ctx), + ) + worker = _worker(is_mtp_eagle, fast_paths) + tokens, _ = worker._forward_linear_draft_loop( + inputs, + md, + spec_metadata, + draft_model, + None, + num_contexts, + b, + num_acc, + None, + accepted_tokens=accepted, + pre_draft_seq_lens=md._seq_lens_cuda.clone(), + ) + torch.cuda.synchronize() + return tokens, md.kv_lens_cuda.clone(), worker._step0_row_pick._picks_rows + + +def _compare(ref_layer, ref, layer, out, b, n_rows): + ref_tokens, ref_kv, ref_picked = ref + tokens, kv, picked = out + assert not ref_picked + assert torch.equal(tokens, ref_tokens) + assert torch.equal(kv, ref_kv) + # The draft KV of every row is written on every step. + assert len(layer.kv_writes) == len(ref_layer.kv_writes) + for got, want in zip(layer.kv_writes, ref_layer.kv_writes): + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-5) + assert layer.kv_writes[0].shape[0] == n_rows + # Step 1's carried hidden states. + for got, want in zip(layer.inputs[1:], ref_layer.inputs[1:]): + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-5) + assert ref_layer.post_rows[0] == n_rows + assert all(r == b for r in ref_layer.post_rows[1:]) + return picked + + +@pytest.mark.parametrize("fast_paths", [False, True]) +# Mixed contexts and generations; contexts only with a single draft step. +@pytest.mark.parametrize("num_contexts,num_gens,k", [(2, 3, 3), (3, 0, 1)]) +def test_mtp_eagle_step0_pick_matches_full(fast_paths, num_contexts, num_gens, k): + b = num_contexts + num_gens + ref_layer = _PlainMtpLayer(seed=7) + ref = _run_loop(ref_layer, True, fast_paths, num_contexts, num_gens, k, 3) + layer = _PickingMtpLayer(seed=7) + out = _run_loop(layer, True, fast_paths, num_contexts, num_gens, k, 3) + n_rows = ref_layer.kv_writes[0].shape[0] + assert _compare(ref_layer, ref, layer, out, b, n_rows) + # The post-attention part ran on one row per request on every step. + assert layer.post_rows == [b] * k + + +def test_eagle3_step0_pick_matches_full(): + num_contexts, num_gens, k = 0, 4, 2 + b = num_contexts + num_gens + ref_layer = _Eagle3Model(seed=11, supports=False) + ref = _run_loop(ref_layer, False, False, num_contexts, num_gens, k, 5) + layer = _Eagle3Model(seed=11, supports=True) + out = _run_loop(layer, False, False, num_contexts, num_gens, k, 5) + n_rows = ref_layer.kv_writes[0].shape[0] + assert _compare(ref_layer, ref, layer, out, b, n_rows) + assert layer.post_rows == [b] * k + + +@pytest.mark.parametrize("layer_cls", [_PlainMtpLayer, _OverridingMtpLayer]) +def test_layer_without_pick_support_is_unaffected(layer_cls): + num_contexts, num_gens, k = 2, 3, 3 + ref_layer = _PlainMtpLayer(seed=7) + ref = _run_loop(ref_layer, True, True, num_contexts, num_gens, k, 3) + layer = layer_cls(seed=7) + out = _run_loop(layer, True, True, num_contexts, num_gens, k, 3) + n_rows = ref_layer.kv_writes[0].shape[0] + assert not _compare(ref_layer, ref, layer, out, num_contexts + num_gens, n_rows) + assert layer.post_rows == ref_layer.post_rows + + +def _mtp_worker(mapping=None): + """An ``MTPWorker`` whose steps around the draft loop are stubbed out.""" + w = object.__new__(MTPWorker) + w.guided_decoder = None + w.sa_enhancer = None + w.mapping = mapping + w._auxiliary_state_handlers = [] + w._row_pick = DraftRowPick("test MTP layers") + w._execute_guided_decoder_if_present = lambda logits: None + w._commit_target_mamba_states = lambda *args: None + w.update_mtp_hidden_states = lambda **kwargs: None + w.change_attn_metadata = lambda *args: None + w.get_draft_kv_cache_manager = lambda resource_manager: None + w.draft_kv_cache_context = lambda md, mgr: contextlib.nullcontext(md) + w._restore_attn_metadata_from_spec_dec = lambda md: None + w._prepare_next_new_tokens = lambda *args: None + w.sample_draft_tokens = lambda logits, spec_metadata, batch_size, draft_step: logits.argmax( + -1 + ).to(torch.int32) + return w + + +def _run_mtp_worker(layer_cls, mapping=None, num_contexts=2, num_gens=3, k=3, seed=9): + """``MTPWorker.forward``'s draft loop over ``k`` MTP layers, in the + layout after ``change_attn_metadata`` (``k`` rows per generation).""" + g = torch.Generator().manual_seed(seed) + seq_lens = torch.randint(2, 9, (num_contexts,), generator=g).tolist() + [k] * num_gens + b, n = num_contexts + num_gens, sum(seq_lens) + md = _Meta(seq_lens, num_contexts, [60] * b) + layers = [layer_cls(seed=100 + i) for i in range(k)] + draft_inputs = { + "input_ids": torch.randint(0, VOCAB, (n,), generator=g, dtype=torch.int32).cuda(), + "position_ids": torch.randint(0, 1000, (n,), generator=g, dtype=torch.int32).cuda(), + "hidden_states": torch.randn(n, H, generator=g).cuda(), + "attn_metadata": md, + } + accepted = torch.randint(0, VOCAB, (b, k + 1), generator=g, dtype=torch.int32).cuda() + num_acc = torch.ones(b, dtype=torch.int32, device="cuda") + worker = _mtp_worker(mapping) + worker.sample_and_accept_draft_tokens = lambda *args: (accepted, num_acc) + worker.prepare_drafter_inputs = lambda **kwargs: draft_inputs + out = worker._forward_impl( + draft_inputs["input_ids"], + draft_inputs["position_ids"], + draft_inputs["hidden_states"], + torch.zeros(b, VOCAB, device="cuda"), + md, + SimpleNamespace(runtime_draft_len=k, batch_indices_cuda=None), + SimpleNamespace(mtp_layers=layers, embed_tokens=None, lm_head=None), + ) + torch.cuda.synchronize() + return out["next_draft_tokens"], layers, worker._row_pick._picks_rows, b, n + + +def test_mtp_worker_layers_pick_matches_full(): + """Vanilla MTP: every MTP layer runs after attention only on each + request's last row; draft tokens and every layer's carried hidden input + match, and each layer's attention still sees (and writes) every row.""" + ref_tokens, ref_layers, ref_picked, b, n = _run_mtp_worker(_PlainMtpLayer) + tokens, layers, picked, _, _ = _run_mtp_worker(_PickingMtpLayer) + assert picked and not ref_picked + assert torch.equal(tokens, ref_tokens) + for layer, ref_layer in zip(layers, ref_layers): + assert layer.post_rows == [b] and ref_layer.post_rows == [n] + assert [x.shape[0] for x in layer.kv_writes] == [n] + torch.testing.assert_close(layer.inputs[0], ref_layer.inputs[0], rtol=1e-5, atol=1e-5) + torch.testing.assert_close(layer.kv_writes[0], ref_layer.kv_writes[0], rtol=1e-5, atol=1e-5) + + +def test_step0_pick_under_attention_dp(): + """Attention DP. The draft loop: a picked step 0 hands the layer the + per-request counts (``subseq_all_rank_num_tokens``, the counts of its + post-attention rows), an unpicked one the full step-0 counts, later steps + always the per-request counts; picking without the per-request counts + fails. Vanilla MTP declines the pick and runs as before.""" + shape = num_contexts, num_gens, k = 2, 3, 3 + b = num_contexts + num_gens + full_counts, per_request_counts = [101, 202], [5, 6] + for layer, is_mtp_eagle, picking in [ + (_PickingMtpLayer(seed=7), True, True), + (_Eagle3Model(seed=11, supports=False), False, False), + ]: + _, _, picked = _run_loop( + layer, is_mtp_eagle, False, *shape, 3, full_counts, per_request_counts + ) + assert picked == picking + step0 = per_request_counts if picking else full_counts + assert layer.all_rank_num_tokens == [step0] + [per_request_counts] * (k - 1) + assert layer.post_rows[0] == (b if picking else layer.kv_writes[0].shape[0]) + with pytest.raises(AssertionError): + _run_loop(_PickingMtpLayer(seed=7), True, False, *shape, 3, full_counts) + + adp = SimpleNamespace(enable_attention_dp=True, cp_size=1) + tokens, layers, picked, _, n = _run_mtp_worker(_PickingMtpLayer, mapping=adp) + ref_tokens, _, _, _, _ = _run_mtp_worker(_PlainMtpLayer) + assert not picked + assert all(layer.post_rows == [n] for layer in layers) + assert torch.equal(tokens, ref_tokens) + + +def test_draft_row_pick_gates(monkeypatch): + """``draft_layer_picks_rows`` and the ``DraftRowPick`` decision, which is + resolved again whenever the draft layer classes change.""" + monkeypatch.delenv("TLLM_SPEC_DISABLE_DRAFT_ROW_PICK", raising=False) + assert draft_layer_picks_rows(_PickingMtpLayer(0)) + assert not draft_layer_picks_rows(None) + assert not draft_layer_picks_rows(_PlainMtpLayer(0)) + assert not draft_layer_picks_rows(_OverridingMtpLayer(0)) + assert not draft_layer_picks_rows(_Eagle3Model(0, supports=False)) + assert draft_layer_picks_rows(_Eagle3Model(0, supports=True)) + assert not draft_layer_picks_rows(_PickingMtpLayer(0), SimpleNamespace(cp_size=2)) + assert draft_layer_picks_rows(_PickingMtpLayer(0), SimpleNamespace(cp_size=1)) + + resolver = DraftRowPick("test") + assert resolver.resolve([_PickingMtpLayer(0)]) + assert not resolver.resolve([_OverridingMtpLayer(0)]) + assert resolver.resolve([_PickingMtpLayer(0)]) + assert not resolver.resolve([_PickingMtpLayer(0), _PlainMtpLayer(0)]) + assert not DraftRowPick("test").resolve([_PickingMtpLayer(0)], decline="attention DP") + + monkeypatch.setenv("TLLM_SPEC_DISABLE_DRAFT_ROW_PICK", "1") + assert not draft_layer_picks_rows(_PickingMtpLayer(0)) + + +class _StubNorm: + """RMSNorm (with the optional residual add) without weights.""" + + weight = None + variance_epsilon = 1e-6 + + def __call__(self, x, residual=None): + if residual is not None: + x = x + residual + residual = x + y = x / x.float().pow(2).mean(-1, keepdim=True).add(1e-6).sqrt() + return y if residual is None else (y, residual) + + +_stub_norm = _StubNorm() + + +def _mtp_layer_cls(name): + if name == "DeepseekV3MTP": + from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3MTP + + return DeepseekV3MTP + from tensorrt_llm._torch.models.modeling_glm import Glm4MTP + + return Glm4MTP + + +@pytest.mark.parametrize( + "layer_name,fused_allreduce", [("DeepseekV3MTP", False), ("Glm4MTP", True)] +) +def test_mtp_layer_forward_pick_rows(layer_name, fused_allreduce): + """The MTP layer's ``forward`` with ``pick_rows`` equals the picked rows of + the full forward; attention sees every row, the MoE only the picked.""" + mtp_cls = _mtp_layer_cls(layer_name) + assert draft_layer_picks_rows(object.__new__(mtp_cls)) + g = torch.Generator().manual_seed(0) + w_eh = (torch.randn(2 * H, H, generator=g) / H**0.5).cuda() + w_moe = (torch.randn(H, H, generator=g) / H**0.5).cuda() + emb = torch.randn(VOCAB, H, generator=g).cuda() + seen = {"attn": [], "moe": []} + + def self_attn(position_ids, hidden_states, attn_metadata, all_reduce_params, **kwargs): + seen["attn"].append(hidden_states.shape[0]) + n = torch.arange(1, hidden_states.shape[0] + 1, device="cuda").unsqueeze(1) + return hidden_states.cumsum(0) / n + + def mlp(hidden_states, all_rank_num_tokens, final_all_reduce_params, **kwargs): + seen["moe"].append(hidden_states.shape[0]) + return torch.relu(hidden_states @ w_moe) + + def allreduce(hidden_states, all_reduce_params): + p = all_reduce_params + return _stub_norm(hidden_states, p.residual) + + from tensorrt_llm._torch.utils import EventType + + shared_head = SimpleNamespace(norm=_stub_norm) + layer = SimpleNamespace( + enorm=_stub_norm, + hnorm=_stub_norm, + event_dict={EventType.Main: torch.cuda.Event(), EventType.MoeShared: torch.cuda.Event()}, + aux_stream=None, + model_config=SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, tp_rank=0, enable_attention_dp=False) + ), + mapping=SimpleNamespace(tp_size=1), + eh_proj=lambda x: x @ w_eh, + input_layernorm=_stub_norm, + self_attn=self_attn, + disable_attn_allreduce=False, + fusion_config=SimpleNamespace( + PRE_MOE_FUSION=fused_allreduce, POST_MOE_FUSION=fused_allreduce + ), + allreduce=allreduce, + post_attention_layernorm=_stub_norm, + mlp=mlp, + shared_head=shared_head, + ) + n = 23 + input_ids = torch.randint(0, VOCAB, (n,), generator=g).cuda() + hidden = torch.randn(n, H, generator=g).cuda() + pick = torch.tensor([4, 9, 15, 22], device="cuda") + kwargs = dict( + input_ids=input_ids, + position_ids=None, + hidden_states=hidden, + embed_tokens=lambda ids: emb[ids], + attn_metadata=None, + ) + full = mtp_cls.forward(layer, **kwargs) + picked = mtp_cls.forward(layer, pick_rows=pick, **kwargs) + torch.testing.assert_close(picked, full[pick], rtol=1e-5, atol=1e-5) + assert seen == {"attn": [n, n], "moe": [n, pick.numel()]} + + +def test_eagle3_decoder_layer_pick_rows(): + from tensorrt_llm._torch.models.modeling_speculative import Eagle3DecoderLayer, Eagle3DraftModel + + g = torch.Generator().manual_seed(1) + w_mlp = (torch.randn(H, H, generator=g) / H**0.5).cuda() + seen = {"attn": [], "mlp": []} + + def self_attn(position_ids, hidden_states, attn_metadata): + seen["attn"].append(hidden_states.shape[0]) + n = torch.arange(1, hidden_states.shape[0] + 1, device="cuda").unsqueeze(1) + return (hidden_states.cumsum(0) / n)[:, :H] + + def mlp(hidden_states): + seen["mlp"].append(hidden_states.shape[0]) + return torch.relu(hidden_states @ w_mlp) + + decoder = SimpleNamespace( + layer_idx=40, + hidden_norm=_stub_norm, + _next_layer_regular=False, + input_layernorm=_stub_norm, + self_attn=self_attn, + post_attention_layernorm=_stub_norm, + mlp=mlp, + ) + model = SimpleNamespace( + supports_pick_rows=True, + embed_tokens=None, + dtype=torch.float32, + _eh_proj_before_attn=False, + num_layers=1, + midlayer=lambda **kw: Eagle3DecoderLayer.forward(decoder, **kw), + norm=_stub_norm, + _return_hidden_post_norm=False, + ) + spec_metadata = SimpleNamespace(maybe_capture_hidden_states=lambda *a, **k: None) + n = 17 + pick = torch.tensor([0, 5, 16], device="cuda") + kwargs = dict( + attn_metadata=SimpleNamespace(all_rank_num_tokens=None), + inputs_embeds=torch.randn(n, H, generator=g).cuda(), + position_ids=None, + spec_metadata=spec_metadata, + hidden_states=torch.randn(n, H, generator=g).cuda(), + ) + full, full_save = Eagle3DraftModel.forward(model, **kwargs) + picked, picked_save = Eagle3DraftModel.forward(model, pick_rows=pick, **kwargs) + torch.testing.assert_close(picked, full[pick], rtol=1e-5, atol=1e-5) + torch.testing.assert_close(picked_save, full_save[pick], rtol=1e-5, atol=1e-5) + assert seen == {"attn": [n, n], "mlp": [n, pick.numel()]} + model.supports_pick_rows = False + with pytest.raises(AssertionError): + Eagle3DraftModel.forward(model, pick_rows=pick, **kwargs) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 46a01a17aabb..0f27aba53e4b 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -98,6 +98,7 @@ def run_draft_forward(*args, **kwargs): is_mtp_eagle=True, draft_kv_cache_context=draft_context, _run_draft_forward=run_draft_forward, + _step0_draft_layer_picks_rows=lambda draft_model: False, ) from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker diff --git a/tests/unittest/_torch/speculative/test_fused_accept.py b/tests/unittest/_torch/speculative/test_fused_accept.py index e4b5bfb37b89..f07fb1dc2e6e 100644 --- a/tests/unittest/_torch/speculative/test_fused_accept.py +++ b/tests/unittest/_torch/speculative/test_fused_accept.py @@ -91,7 +91,7 @@ def test_fused_accept_matches_unfused(force, num_contexts, num_gens, monkeypatch ) def test_step0_epilogue_matches_stock(num_contexts, num_gens, k, rows): # rows = 0: 1-D position ids, else [rows, tokens] MRoPE positions. - from tensorrt_llm._torch.speculative.fused_accept import step0_epilogue + from tensorrt_llm._torch.speculative.fused_accept import step0_epilogue, step0_gather_ids g = torch.Generator(device="cuda").manual_seed(num_contexts * 1000 + num_gens * 10 + k) b, H = num_contexts + num_gens, 4096 @@ -145,6 +145,31 @@ def test_step0_epilogue_matches_stock(num_contexts, num_gens, k, rows): ) is None ) + # step0_gather_ids gives the same ids before the draft forward; the epilogue of a draft layer that kept + # only those rows (picked_gather_ids) reuses the ids and rows instead of gathering, with the same results. + ids = step0_gather_ids(seq_lens, num_acc, num_contexts, b, k, num_ctx_tokens) + assert ids is not None and ids.dtype == torch.int64 and torch.equal(ids, ids_ref) + hid_picked = hidden[ids] + seq_picked, kv_picked = torch.full_like(seq_live, 7), kv.clone() + picked = step0_epilogue( + seq_lens, + num_acc, + hid_picked, + pos, + seq_picked, + kv_picked, + num_contexts, + b, + k, + num_ctx_tokens, + picked_gather_ids=ids, + ) + torch.cuda.synchronize() + assert picked is not None and picked.gather_ids is ids and picked.hidden is hid_picked + for j in range(1, k): + assert torch.equal(picked.positions(j), pos_ref[j - 1]), j + assert torch.equal(seq_picked, seq_live) and torch.equal(kv_picked, kv_ref) + assert step0_gather_ids(seq_lens, num_acc, num_contexts, 2048, k, num_ctx_tokens) is None # Draft tail: row 0 of the token buffer = stock _prepare_next_new_tokens column 0. acc_tok = torch.randint(0, 248320, (b, k + 3), generator=g, device="cuda", dtype=torch.int32)[ :, : k + 2