do not remove rejected tokens in prefill to remove sync op

Author: RunningLeon

lmdeploy/pytorch/engine/model_agent.py

@@ -391,7 +391,8 @@ async def async_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_ou
         await asyncio.sleep(0)
         return output
 
-    async def _async_model_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap):
+    async def _async_model_forward(self, inputs: ModelInputs, spec_inputs: SpecDecodeInputs, swap_in_map: SwapMap,
+                                   swap_out_map: SwapMap):
         """Model forward.
 
         Args:
@@ -433,11 +434,11 @@ async def __long_context_single_forward(new_inputs):
             outputs = await __forward(inputs)
 
         loop_count = self.num_spec_tokens - 1
-        draft_token_ids, model_metas, target_hidden_states = self.proposer.get_outputs(outputs, inputs)
+        draft_token_ids, model_metas, target_hidden_states = self.proposer.get_outputs(outputs, inputs, spec_inputs)
         draft_tokens_li = [draft_token_ids]
         if loop_count > 0:
-            inputs = self.proposer.update_inputs_decoding(inputs, draft_token_ids.transpose(0, 1), target_hidden_states,
-                                                          model_metas)
+            inputs = self.proposer.update_inputs_decoding(inputs, spec_inputs, draft_token_ids.transpose(0, 1),
+                                                          target_hidden_states, model_metas)
             for loop_idx in range(loop_count):
                 outputs = await self.async_forward(inputs, swap_in_map=dict(), swap_out_map=dict())
                 draft_token_ids, model_metas, target_hidden_states = self.proposer.get_outputs(outputs, inputs)
@@ -466,14 +467,17 @@ async def async_model_forward(self,
                 spec_inputs.num_rejected_tokens = num_rejected_tokens
                 spec_inputs.reject_sample_tokens = output_token_ids
                 spec_inputs.next_token_ids = last_token_ids
+                spec_inputs.last_token_indices = model_inputs.seq_length.cumsum(0) - 1 - num_rejected_tokens
             else:
                 spec_inputs.next_token_ids = spec_inputs.bonus_token_ids
                 output_token_ids = spec_inputs.next_token_ids.unsqueeze(-1)
+                spec_inputs.last_token_indices = model_inputs.seq_length.cumsum(0) - 1
 
             with record_function('draft_prepare_inputs'):
                 draft_model_inputs = self.proposer.prepare_inputs(model_inputs, spec_inputs)
 
             new_draft_tokens = await self._async_model_forward(draft_model_inputs,
+                                                               spec_inputs,
                                                                swap_in_map=swap_in_map,
                                                                swap_out_map=swap_out_map)
             outputs = dict(output_token_ids=output_token_ids, spec_token_ids=new_draft_tokens)

lmdeploy/pytorch/model_inputs.py

@@ -132,6 +132,7 @@ class SpecDecodeInputs:
     next_token_ids: torch.LongTensor = None
     num_rejected_tokens: torch.LongTensor = None
     reject_sample_tokens: torch.LongTensor = None
+    last_token_indices: torch.LongTensor = None
 
 
 @dataclass

lmdeploy/pytorch/spec_decode/base.py

@@ -80,7 +80,10 @@ def build_model(self, empty_init: bool, target_model: torch.nn.Module = None):
         self.model = patched_model
         self.target_model = target_model
 
-    def get_outputs(self, model_outputs: Dict[str, torch.Tensor], model_inputs: ModelInputs):
+    def get_outputs(self,
+                    model_outputs: Dict[str, torch.Tensor],
+                    model_inputs: ModelInputs,
+                    spec_inputs: SpecDecodeInputs = None):
         """Get outputs."""
         raise NotImplementedError()
 
@@ -97,20 +100,24 @@ def _forward(self, model_inputs: ModelInputs, cache_engine: CacheEngine = None,
                                    cache_engine=cache_engine,
                                    stream=stream)
 
-    def update_inputs_decoding(self, model_inputs: ModelInputs, input_ids: torch.Tensor,
-                               target_hidden_states: torch.Tensor, model_metas: List[Any]):
+    def update_inputs_decoding(self, model_inputs: ModelInputs, spec_inputs: SpecDecodeInputs,
+                               next_input_ids: torch.Tensor, target_hidden_states: torch.Tensor,
+                               model_metas: List[Any]):
         """Update to decoding inputs."""
         model_inputs.is_decoding = True
         batch_size = model_inputs.seq_length.size(0)
-        model_inputs.input_ids = input_ids
+        model_inputs.input_ids = next_input_ids
         model_inputs.max_q_seqlen = 1
         model_inputs.max_kv_seqlen += 1
         model_inputs.sum_kv_seqlen += model_inputs.seq_length.numel()
         model_inputs.history_lengths += model_inputs.seq_length
+        if spec_inputs.num_rejected_tokens is not None:
+            model_inputs.history_lengths -= spec_inputs.num_rejected_tokens
         model_inputs.seq_length = model_inputs.seq_length.new_ones(batch_size)
         model_inputs.target_position_ids = model_inputs.history_lengths.unsqueeze(0).clone()
         model_inputs.model_metas = model_metas
         model_inputs.target_hidden_states = target_hidden_states
+        model_inputs.spec_metadata = None
         return model_inputs
 
     @record_function('draft_get_logits')

lmdeploy/pytorch/spec_decode/deepseek_mtp.py

@@ -2,8 +2,6 @@
 from typing import Dict
 
 import torch
-import triton
-import triton.language as tl
 
 from lmdeploy.utils import get_logger
 
@@ -16,54 +14,35 @@
 @SPEC_PROPOSERS.register_module(name='deepseek_mtp')
 class DeepseekMTP(BaseSpecProposer):
 
-    def get_outputs(self, model_outputs: Dict[str, torch.Tensor], model_inputs: ModelInputs):
+    def get_outputs(self,
+                    model_outputs: Dict[str, torch.Tensor],
+                    model_inputs: ModelInputs,
+                    spec_inputs: SpecDecodeInputs = None):
         """Get outputs."""
         hidden_states = model_outputs['hidden_states']
         model_metas = model_outputs['model_metas']
         if not model_inputs.is_decoding:
-            if model_inputs.seq_length.size(0) == 1:
+            assert spec_inputs is not None, 'spec_inputs should be provided for prefill mode'
+            if model_inputs.seq_length.size(0) == 1 and spec_inputs.num_rejected_tokens is None:
                 hidden_states = hidden_states[:, -1:]
             else:
-                last_token_loc = model_inputs.seq_length.cumsum(0) - 1
+                last_token_loc = spec_inputs.last_token_indices
                 hidden_states = hidden_states[:, last_token_loc]
+
         logits = self.get_logits(hidden_states)[0]
         draft_token_ids = logits.argmax(dim=-1, keepdim=True)
         return draft_token_ids, model_metas, hidden_states
 
     def prepare_inputs(self, model_inputs: ModelInputs, spec_inputs: SpecDecodeInputs):
         """Prepare inputs."""
         spec_metadata = model_inputs.spec_metadata
-
-        if spec_metadata.draft_token_ids is None:
-            input_ids = model_inputs.input_ids
-            seq_length = model_inputs.seq_length
-        else:
-            # select input ids
-            query_lens = model_inputs.seq_length
-            batch_size = query_lens.size(0)
-            cum_query_lens = query_lens.new_zeros((batch_size + 1), dtype=torch.long)
-            cum_qery_lens_new = query_lens.new_zeros((batch_size + 1), dtype=torch.long)
-            torch.cumsum(query_lens, dim=0, out=cum_query_lens[1:])
-            query_lens_new = query_lens - spec_inputs.num_rejected_tokens
-            torch.cumsum(query_lens_new, dim=0, out=cum_qery_lens_new[1:])
-            keep_token_indices = query_lens.new_zeros(
-                model_inputs.input_ids.size(1) - spec_inputs.num_rejected_tokens.sum())
-            cal_token_indices[(batch_size, )](keep_token_indices, cum_query_lens, cum_qery_lens_new, BLOCK_SIZE=1024)
-            input_ids = model_inputs.input_ids[:, keep_token_indices]
-            seq_length = query_lens_new
-
-            spec_inputs.target_hidden_states = spec_inputs.target_hidden_states[:, keep_token_indices]
-            if spec_inputs.target_position_ids is not None:
-                spec_inputs.target_position_ids = spec_inputs.target_position_ids[:, keep_token_indices]
-
-        # offset by 1 token
+        input_ids = model_inputs.input_ids
+        seq_length = model_inputs.seq_length
+        last_token_indices = spec_inputs.last_token_indices
+        # # offset by 1 token
         input_ids[:, :-1] = input_ids[:, 1:].clone()
-        # update next tokens
-        if seq_length.size(0) == 1:
-            input_ids[:, -1:] = spec_inputs.next_token_ids
-        else:
-            last_token_indices = seq_length.cumsum(0) - 1
-            input_ids[:, last_token_indices] = spec_inputs.next_token_ids
+        # # update next tokens
+        input_ids[:, last_token_indices] = spec_inputs.next_token_ids
         # use new inputs
         return ModelInputs(
             input_ids=input_ids,
@@ -77,30 +56,5 @@ def prepare_inputs(self, model_inputs: ModelInputs, spec_inputs: SpecDecodeInput
             is_decoding=model_inputs.is_decoding,
             target_hidden_states=spec_inputs.target_hidden_states,
             target_position_ids=spec_inputs.target_position_ids,
-        )
-
-
-@triton.jit
-def cal_token_indices(
-    token_indices_ptr,
-    cum_query_lens_ptr,
-    cum_new_query_lens_ptr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Calculate the token indices based on rejection sampler results."""
-    pid = tl.program_id(0)
-
-    start_pos = tl.load(cum_new_query_lens_ptr + pid)
-    end_pos = tl.load(cum_new_query_lens_ptr + pid + 1)
-    num_tokens = end_pos - start_pos
-
-    index_start = tl.load(cum_query_lens_ptr + pid)
-
-    num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE)
-    for i in tl.range(num_blocks):
-        offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-        tl.store(
-            token_indices_ptr + start_pos + offset,
-            index_start + offset,
-            mask=offset < num_tokens,
+            spec_metadata=spec_metadata,
         )

lmdeploy/pytorch/spec_decode/eagle3.py

@@ -6,7 +6,7 @@
 from lmdeploy.utils import get_logger
 
 from ..config import ModelConfig
-from ..model_inputs import ModelInputs
+from ..model_inputs import ModelInputs, SpecDecodeInputs
 from .base import SPEC_PROPOSERS
 from .deepseek_mtp import DeepseekMTP
 
@@ -30,27 +30,26 @@ def get_target_hidden_size(self, model_config: ModelConfig):
         hidden_size = getattr(hf_config, 'target_hidden_size', hf_config.hidden_size)
         return hidden_size * 3
 
-    def get_outputs(self, model_outputs: Dict[str, torch.Tensor], model_inputs: ModelInputs):
+    def get_outputs(self,
+                    model_outputs: Dict[str, torch.Tensor],
+                    model_inputs: ModelInputs,
+                    spec_inputs: SpecDecodeInputs = None):
         """Get outputs."""
         hidden_states = model_outputs['hidden_states']
         hidden_states_prenorm = model_outputs['hidden_states_prenorm']
         model_metas = model_outputs['model_metas']
         if not model_inputs.is_decoding:
-            if model_inputs.seq_length.size(0) == 1:
+            assert spec_inputs is not None, 'spec_inputs should be provided for prefill mode'
+            if model_inputs.seq_length.size(0) == 1 and spec_inputs.num_rejected_tokens is None:
                 hidden_states = hidden_states[:, -1:]
                 hidden_states_prenorm = hidden_states_prenorm[:, -1:]
             else:
-                last_token_loc = model_inputs.seq_length.cumsum(0) - 1
+                last_token_loc = spec_inputs.last_token_indices
                 hidden_states = hidden_states[:, last_token_loc]
                 hidden_states_prenorm = hidden_states_prenorm[:, last_token_loc]
 
         logits = self.get_logits(hidden_states)[0]
         draft_token_ids = logits.argmax(dim=-1, keepdim=True)
-        device = draft_token_ids.device
-        dtype = draft_token_ids.dtype
         # token mapping
-        if self.draft_id_to_target_id.device != device or self.draft_id_to_target_id.dtype != dtype:
-            self.draft_id_to_target_id = self.draft_id_to_target_id.to(dtype=draft_token_ids.dtype,
-                                                                       device=draft_token_ids.device)
         draft_token_ids = self.draft_id_to_target_id[draft_token_ids]
         return draft_token_ids, model_metas, hidden_states_prenorm