Spec decode warmup support

Signed-off-by: Chen Haifeng <haifeng.chen@intel.com>

Author: jerrychenhf

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -3410,7 +3410,9 @@ def sample_tokens(self, grammar_output: "GrammarOutput | None") -> ModelRunnerOu
         for req_id in self.input_batch.req_ids[:num_reqs]:
             req_state = self.requests[req_id]
             i = self.input_batch.req_id_to_index[req_id]
-            seq_len = (req_state.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id])
+            # Cannot use num_computed_tokens + num_scheduled_tokens here
+            # as it may include rejected spec decode tokens
+            seq_len = self.input_batch.num_tokens_no_spec[i]
             token_ids = postprocessed_sampled_token_ids[i]
             num_tokens = len(token_ids)
             self.input_batch.token_ids_cpu[i, seq_len:seq_len + num_tokens] = token_ids
@@ -3996,7 +3998,16 @@ def _add_dummy_request(self,
                            scheduled_tokens,
                            is_prompt,
                            block_id=0):
-        num_blocks = round_up(total_tokens, self.block_size) // self.block_size
+        # Spec decode: blocks should include look ahead tokens (eagle)
+        total_tokens_for_blocks = total_tokens
+        if self.speculative_config and self.speculative_config.use_eagle():
+            # Consider the block space for draft tokens to propose
+            total_tokens_for_blocks += self.speculative_config.num_speculative_tokens
+            # Check the limit of the max model length
+            if total_tokens_for_blocks > self.max_model_len:
+                total_tokens_for_blocks = self.max_model_len
+
+        num_blocks = round_up(total_tokens_for_blocks, self.block_size) // self.block_size
         prompt_token_ids = list(range(total_tokens))
 
         req_id = f'{len(requests)}'
@@ -4075,14 +4086,20 @@ def _add_dummy_unified_request(self, requests, is_prompt, is_unique, block_num,
         requests.append(req)
         scheduled_tokens[req_id] = num_scheduled_tokens
 
-    @staticmethod
-    def _generate_seq_lengths(num_samples, num_blocks, block_size):
+    def _generate_seq_lengths(self, num_samples, num_blocks, block_size):
         assert num_samples <= num_blocks
         blocks = [num_blocks // num_samples] * num_samples
         missing_blocks = num_blocks - sum(blocks)
         for i in range(missing_blocks):
             blocks[i] += 1
-        seq_lengths = [b * block_size - 1 for b in blocks]
+
+        # Leave space for the output token and draft tokens to propose
+        num_lookahead_tokens = 1
+        if self.speculative_config and self.speculative_config.use_eagle():
+            # Consider the token space for draft tokens to propose
+            # The draft tokens for eagle consumes block table space
+            num_lookahead_tokens += self.speculative_config.num_speculative_tokens
+        seq_lengths = [b * block_size - num_lookahead_tokens for b in blocks]
         return seq_lengths
 
     def distribute_sum_evenly(self, total_sum, max_length):
@@ -4222,6 +4239,12 @@ def _prepare_dummy_scenario(self, prompt_cfg, decode_cfg):
                                                      prompt_num_blocks)
             for _ in range(prompt_bs):
                 for tokens, context_len in zip(prompt_total_tokens, prompt_num_context_blocks):
+                    if self.speculative_config and self.speculative_config.use_eagle():
+                        # Leave the block space for draft tokens to propose
+                        # The draft tokens for eagle consumes block table space
+                        num_speculative_tokens = self.speculative_config.num_speculative_tokens
+                        tokens -= num_speculative_tokens
+                        prompt_query_len -= num_speculative_tokens
                     self._add_dummy_request(requests,
                                             scheduled_tokens,
                                             num_computed_tokens=(context_len * self.block_size),