From 889f641aed79508d8787b7369e1f549bc96ee9d4 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Mon, 16 Jun 2025 05:12:07 -0700 Subject: [PATCH 1/4] fix the deterministic issue in MTP eagle path. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../_torch/models/modeling_deepseekv3.py | 32 ++--- tensorrt_llm/_torch/speculative/mtp.py | 118 +++++++++--------- 2 files changed, 74 insertions(+), 76 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index d0b00450f3b5..518f9a76b418 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -134,19 +134,23 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): eps=config.rms_norm_eps, dtype=config.torch_dtype) - def forward(self, hidden_states: torch.Tensor, lm_head: Linear, - attn_metadata: AttentionMetadata) -> torch.Tensor: - if attn_metadata is not None: - last_tokens = torch.cumsum( - attn_metadata.seq_lens_cuda, - dim=0, - dtype=torch.long, - ) - 1 - last_token_hidden_states = hidden_states[last_tokens] - else: - last_token_hidden_states = hidden_states[-1].unsqueeze(0) + def forward(self, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: AttentionMetadata, + return_context_logits: bool = False) -> torch.Tensor: + if not return_context_logits: + if attn_metadata is not None: + last_tokens = torch.cumsum( + attn_metadata.seq_lens_cuda, + dim=0, + dtype=torch.long, + ) - 1 + hidden_states = hidden_states[last_tokens] + else: + hidden_states = hidden_states[-1].unsqueeze(0) - logits = lm_head(last_token_hidden_states) + logits = lm_head(hidden_states) return logits @@ -1039,9 +1043,7 @@ def forward( else: hidden_states, _ = self.shared_head.norm(hidden_states, residual) - logits = self.shared_head(hidden_states, lm_head, attn_metadata).float() - - return hidden_states, logits + return hidden_states class DeepseekV3Model(DecoderModel): diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 82ed94ab0fd6..a004ade3b9c3 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -375,9 +375,11 @@ def forward( num_accepted_tokens=num_accepted_tokens, spec_metadata=spec_metadata, attn_metadata=attn_metadata) - hidden_states, logits = mtp_layer(lm_head=lm_head, - embed_tokens=embed_tokens, - **draft_inputs) + hidden_states = mtp_layer(lm_head=lm_head, + embed_tokens=embed_tokens, + **draft_inputs) + logits = mtp_layer.shared_head(hidden_states, lm_head, + attn_metadata).float() previous_layer_draft_tokens = self.draft_sampler(logits) next_draft_tokens.append(previous_layer_draft_tokens) @@ -1066,6 +1068,7 @@ def forward( ): batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts # Sample and verify draft tokens raw_logits = logits @@ -1079,21 +1082,35 @@ def forward( # Prepare inputs for the 1st MTP layer position_ids = position_ids.squeeze(0) - inputs = self.prepare_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata) + last_tokens_idx = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + inputs = self.prepare_drafter_inputs(input_ids=input_ids, + position_ids=position_ids, + last_tokens_idx=last_tokens_idx, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata) # Predict draft tokens next_draft_tokens = [] for i in range(self.mtp_num_modules): - hidden_states, logits = mtp_layers[0](lm_head=lm_head, - embed_tokens=embed_tokens, - **inputs) + hidden_states = mtp_layers[0](lm_head=lm_head, + embed_tokens=embed_tokens, + **inputs) + if i == 0: + start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] * + (self.mtp_num_modules + 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) + else: + # All of the seq_len are 1, use batch_indices_cuda as gather_ids + gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + logits = mtp_layers[0].shared_head(hidden_states[gather_ids], + lm_head, attn_metadata, True) new_draft_token = self.draft_sampler(logits) next_draft_tokens.append(new_draft_token) # update inputs @@ -1102,17 +1119,29 @@ def forward( dim=0, dtype=torch.long, ) - 1 - position_ids = inputs["position_ids"][last_tokens] + 1 - hidden_states = hidden_states[last_tokens] - attn_metadata._seq_lens[:attn_metadata.num_contexts].fill_(1) - attn_metadata._seq_lens_cuda[:attn_metadata.num_contexts].fill_(1) - attn_metadata.on_update() - # cannot run generation if their is no kv cache - if inputs["attn_metadata"].kv_cache_manager is not None: - attn_metadata.host_request_types[:attn_metadata. - num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - if i == 0 and num_contexts > 0 and attn_metadata.enable_flash_mla: + hidden_states = hidden_states[gather_ids] + position_ids = inputs["position_ids"][gather_ids] + 1 + # update attn_metadata + if i == 0: + attn_metadata._seq_lens[:attn_metadata.num_contexts].fill_(1) + attn_metadata._seq_lens_cuda[:attn_metadata.num_contexts].fill_( + 1) + attn_metadata.on_update() + # cannot run generation if their is no kv cache + has_kv_cache = inputs[ + "attn_metadata"].kv_cache_manager is not None + if has_kv_cache: + attn_metadata.host_request_types[:attn_metadata. + num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + # update kv_lens_cuda + if hasattr(attn_metadata, 'kv_lens_cuda'): + attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( + self.mtp_num_modules - + num_accepted_tokens[num_contexts:]) + attn_metadata.kv_lens_cuda[:num_contexts] += 1 + # update metadata for flash mla + if has_kv_cache and num_contexts > 0 and attn_metadata.enable_flash_mla: reorder_block_ids_per_seq = torch.cat([ attn_metadata. kv_block_ids_per_seq[num_contexts:batch_size], @@ -1120,7 +1149,7 @@ def forward( ]) attn_metadata.block_ids_per_seq[:batch_size, :].copy_( reorder_block_ids_per_seq, non_blocking=True) - if hasattr(attn_metadata, 'kv_lens_cuda'): + elif hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[:batch_size] += 1 # support attention dp if spec_metadata.all_rank_num_tokens is not None: @@ -1159,22 +1188,15 @@ def prepare_drafter_inputs( self, input_ids: torch.IntTensor, position_ids: torch.IntTensor, + last_tokens_idx: torch.LongTensor, hidden_states: torch.Tensor, accepted_tokens: torch.Tensor, - num_accepted_tokens: torch.Tensor, attn_metadata: AttentionMetadata, spec_metadata: MTPSpecMetadata, ): - batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_ctx_tokens = attn_metadata.num_ctx_tokens - hidden_size = hidden_states.shape[1] - last_tokens_idx = torch.cumsum( - attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 # context - hidden_states_ctx = hidden_states[:attn_metadata.num_ctx_tokens, :] input_ctx_ids = input_ids[:attn_metadata.num_ctx_tokens] input_ids_ctx = torch.empty_like(input_ctx_ids, dtype=torch.int32, @@ -1182,38 +1204,12 @@ def prepare_drafter_inputs( input_ids_ctx[:-1].copy_(input_ctx_ids[1:]) input_ids_ctx[ last_tokens_idx[:num_contexts]] = accepted_tokens[:num_contexts, 0] - position_ids_ctx = position_ids[:num_ctx_tokens] # generation - gen_batch_idx = spec_metadata.batch_indices_cuda[:num_gens] - gen_token_idx = num_accepted_tokens[num_contexts:] - 1 - hidden_states_gen = hidden_states[attn_metadata.num_ctx_tokens:, :] - hidden_states_gen = hidden_states_gen.reshape(num_gens, - self.mtp_num_modules + 1, - hidden_size) - hidden_states_gen = hidden_states_gen[gen_batch_idx, gen_token_idx, :] - accepted_tokens_gen = accepted_tokens[num_contexts:, :] - input_ids_gen = accepted_tokens_gen[gen_batch_idx, gen_token_idx] - position_ids_gen = position_ids[num_ctx_tokens:].reshape( - num_gens, self.mtp_num_modules + 1) - position_ids_gen = position_ids_gen[gen_batch_idx, gen_token_idx] + input_ids_gen = accepted_tokens[num_contexts:, :].flatten() # get draft inputs input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0) - hidden_states = torch.concat([hidden_states_ctx, hidden_states_gen], - dim=0) - position_ids = torch.concat([position_ids_ctx, position_ids_gen], dim=0) - - # change attn_metadata - attn_metadata._seq_lens[num_contexts:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[num_contexts:batch_size].fill_(1) - attn_metadata.on_update() - if hasattr(attn_metadata, 'kv_lens_cuda'): - # Note that it's important to not free the seq_lens_cuda - # buffer once the graph has been captured also - this will invalidate - # the graph and force an expensive recapture. - attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( - self.mtp_num_modules + 1 - num_accepted_tokens[num_contexts:]) return { "input_ids": input_ids, From 3472018adf7f59b1c82260dc1d8f699c4197c81f Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 17 Jun 2025 00:00:03 -0700 Subject: [PATCH 2/4] fix. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/mtp.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index a004ade3b9c3..1cf8da986f18 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1053,6 +1053,7 @@ class MTPEagleWorker(MTPWorker): def __init__(self, spec_config: MTPConfig): super().__init__(spec_config) self.mtp_num_modules = spec_config.num_nextn_predict_layers + self.is_thop = False def forward( self, @@ -1123,9 +1124,8 @@ def forward( position_ids = inputs["position_ids"][gather_ids] + 1 # update attn_metadata if i == 0: - attn_metadata._seq_lens[:attn_metadata.num_contexts].fill_(1) - attn_metadata._seq_lens_cuda[:attn_metadata.num_contexts].fill_( - 1) + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) attn_metadata.on_update() # cannot run generation if their is no kv cache has_kv_cache = inputs[ From fb382994a05a63baa5e6e8099176d491ce7f304f Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 17 Jun 2025 02:57:02 -0700 Subject: [PATCH 3/4] fix attention dp for mtp eagle. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../_torch/models/modeling_deepseekv3.py | 5 +- tensorrt_llm/_torch/speculative/mtp.py | 52 +++++++++++-------- 2 files changed, 33 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 518f9a76b418..ff17f24d2fb5 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -980,10 +980,9 @@ def forward( input_ids: torch.IntTensor, position_ids: torch.IntTensor, hidden_states: torch.Tensor, - lm_head: Linear, embed_tokens: Embedding, attn_metadata: AttentionMetadata, - spec_metadata: MTPSpecMetadata, + all_rank_num_tokens: Optional[List[int]] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -1024,7 +1023,7 @@ def forward( # MoE hidden_states = self.mlp( hidden_states, - all_rank_num_tokens=spec_metadata.all_rank_num_tokens, + all_rank_num_tokens=all_rank_num_tokens, final_all_reduce_params=AllReduceParams( enable_allreduce=not (self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1)), diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 1cf8da986f18..d82dd1c5512e 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -140,6 +140,13 @@ class MTPSpecMetadata(SpecMetadata): slot_ids: Optional[torch.Tensor] = None # The index of the batche inputs batch_indices_cuda: Optional[torch.Tensor] = None + # The number of sequences for speculative model/layer of different rank + _all_rank_num_seqs: Optional[List[int]] = None + # This is used for attention dp in the MTP Eagle worker. The numbers of input + # tokens varies between the 1st draft forward and subsequent ones. To support + # CUDA graph, we use this tensor to store the number of input tokens for the + # subsequence draft forward. + subseq_all_rank_num_tokens: Optional[List[int]] = None def __post_init__(self) -> None: if self.mtp_hidden_states_manager is not None: @@ -166,6 +173,16 @@ def __post_init__(self) -> None: device='cuda', ) + @property + def all_rank_num_seqs(self): + return self._all_rank_num_seqs + + @all_rank_num_seqs.setter + def all_rank_num_seqs(self, value: List[int]): + self._all_rank_num_seqs = value + if self.spec_dec_mode.is_mtp_eagle(): + self.subseq_all_rank_num_tokens = value + def prepare(self): assert self.request_ids is not None num_seqs = len(self.request_ids) @@ -176,10 +193,11 @@ def prepare(self): pin_memory=True) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - # MTP module need different number of input tokens in generation phase - if self.spec_dec_mode.is_mtp_eagle(): - self.num_tokens -= (self.num_generations) * self.mtp_num_modules - else: + # MTP vanilla worker uses total max_draft_tokens input tokens in generation phase, + # while MTP Eagle worker uses (max_draft_tokens + 1) input tokens in the 1st draft + # forward and only one input token in the following draft forward. + # This num_tokens is used to set the all_rank_num_tokens for attention dp. + if not self.spec_dec_mode.is_mtp_eagle(): self.num_tokens -= self.num_generations if self.mtp_hidden_states_manager is not None: # MTP vanilla or use relaxed acceptance @@ -375,9 +393,7 @@ def forward( num_accepted_tokens=num_accepted_tokens, spec_metadata=spec_metadata, attn_metadata=attn_metadata) - hidden_states = mtp_layer(lm_head=lm_head, - embed_tokens=embed_tokens, - **draft_inputs) + hidden_states = mtp_layer(embed_tokens=embed_tokens, **draft_inputs) logits = mtp_layer.shared_head(hidden_states, lm_head, attn_metadata).float() previous_layer_draft_tokens = self.draft_sampler(logits) @@ -1023,7 +1039,6 @@ def prepare_drafter_inputs( "position_ids": position_ids, "hidden_states": return_hidden_states, "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, } def draft_sampler( @@ -1096,10 +1111,11 @@ def forward( # Predict draft tokens next_draft_tokens = [] for i in range(self.mtp_num_modules): - hidden_states = mtp_layers[0](lm_head=lm_head, - embed_tokens=embed_tokens, - **inputs) if i == 0: + hidden_states = mtp_layers[0]( + embed_tokens=embed_tokens, + all_rank_num_tokens=spec_metadata.all_rank_num_tokens, + **inputs) start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] * (self.mtp_num_modules + 1)).long() gather_ids_gen = (start_ids_gen + @@ -1108,6 +1124,10 @@ def forward( gather_ids = torch.concat( [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) else: + hidden_states = mtp_layers[0](embed_tokens=embed_tokens, + all_rank_num_tokens=spec_metadata. + subseq_all_rank_num_tokens, + **inputs) # All of the seq_len are 1, use batch_indices_cuda as gather_ids gather_ids = spec_metadata.batch_indices_cuda[:batch_size] logits = mtp_layers[0].shared_head(hidden_states[gather_ids], @@ -1115,11 +1135,6 @@ def forward( new_draft_token = self.draft_sampler(logits) next_draft_tokens.append(new_draft_token) # update inputs - last_tokens = torch.cumsum( - attn_metadata.seq_lens_cuda, - dim=0, - dtype=torch.long, - ) - 1 hidden_states = hidden_states[gather_ids] position_ids = inputs["position_ids"][gather_ids] + 1 # update attn_metadata @@ -1151,15 +1166,11 @@ def forward( reorder_block_ids_per_seq, non_blocking=True) elif hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[:batch_size] += 1 - # support attention dp - if spec_metadata.all_rank_num_tokens is not None: - spec_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs inputs = { "input_ids": new_draft_token, "position_ids": position_ids, "hidden_states": hidden_states, "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, } next_draft_tokens = torch.stack(next_draft_tokens, dim=1) @@ -1216,5 +1227,4 @@ def prepare_drafter_inputs( "position_ids": position_ids, "hidden_states": hidden_states, "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, } From d6276c0631b7e4ae984c42e2e845fbf3b0e4e85c Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 17 Jun 2025 21:28:43 -0700 Subject: [PATCH 4/4] fix accepted_tokens tensor. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mtpOp.cpp | 4 ++-- tensorrt_llm/_torch/speculative/mtp.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/thop/mtpOp.cpp b/cpp/tensorrt_llm/thop/mtpOp.cpp index dd880f2b7921..d16efedef64d 100644 --- a/cpp/tensorrt_llm/thop/mtpOp.cpp +++ b/cpp/tensorrt_llm/thop/mtpOp.cpp @@ -109,8 +109,8 @@ std::tuple mtp_sampling_and_accepted_draft_tokens_op(th: TLLM_CHECK(draftTokensSizes[0] == (numGenerationRequest * numMTPModules)); auto stream = at::cuda::getCurrentCUDAStream(logits.get_device()); - auto acceptedTokens = torch::empty( - {batchSize, numMTPModules + 1}, at::TensorOptions().dtype(torch::kInt32).device(logits.device())); + auto acceptedTokens + = torch::ones({batchSize, numMTPModules + 1}, at::TensorOptions().dtype(torch::kInt32).device(logits.device())); auto numAcceptedTokens = torch::ones({batchSize}, at::TensorOptions().dtype(torch::kInt32).device(logits.device())); // Fill params diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index d82dd1c5512e..2b6c580411b3 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -745,12 +745,13 @@ def sample_and_accept_draft_tokens( logits = logits.unsqueeze(0) # The return buffer - accepted_tokens = torch.empty((batch_size, (mtp_num_modules + 1)), - dtype=torch.int, - device=logits.device) - num_accepted_tokens = torch.ones(batch_size, + if self.spec_config.use_relaxed_acceptance_for_thinking or not self.is_thop: + accepted_tokens = torch.ones((batch_size, (mtp_num_modules + 1)), dtype=torch.int, device=logits.device) + num_accepted_tokens = torch.ones(batch_size, + dtype=torch.int, + device=logits.device) if self.spec_config.use_relaxed_acceptance_for_thinking: mtp_relaxed_delta_pool = spec_metadata.mtp_hidden_states_manager.mtp_relaxed_delta_pool @@ -1068,7 +1069,6 @@ class MTPEagleWorker(MTPWorker): def __init__(self, spec_config: MTPConfig): super().__init__(spec_config) self.mtp_num_modules = spec_config.num_nextn_predict_layers - self.is_thop = False def forward( self,