Spec decode warmup support (vllm-project#624)

GAUDISW-242931

Because currently spec decode flatten the spec decode tokens into
[batch_size * num_tokens, 1], we can warmup the decode shapes as it was.
The thing changed is the maximum batch_size we should warmup in the
configuration because the real batch size is batch_size * num_tokens
which is num_tokens (1 + num_speculative_tokens) times of original batch
size.

The thing to care in the warmup is the draft token (and block) space for
the proposing process in eagle. We need to leave out the
num_speculative_tokens space to use by propose for eagle.

Other care needs to be taken (already done in the PR of support
num_speculative_tokens > 1) is warmup will be run in compile only mode
without the real computation happening. So the operations for
prepare_attn_metadata in the drafter which depends on the real position
values must be done on CPU)

Another issue of handling no spec decode tokens for decode phase has
already been handled vllm-project#593

---------

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

Author: jerrychenhf

vllm_gaudi/extension/bucketing/common.py

@@ -39,20 +39,30 @@ class HPUBucketingManager():
     prompt_buckets: List[Tuple[int, int, int]] = []
     decode_buckets: List[Tuple[int, int, int]] = []
     unified_buckets: List[Tuple[int, int, int]] = []
+    # Seed buckets are the buckets originally generated from bucketing configuration
+    # Spec decode may automatically add new buckets based on the seed buckets
+    seed_decode_buckets: List[Tuple[int, int, int]] = None
     initialized = False
 
     def __new__(cls, *args, **kwargs):
         if not cls._instance:
             cls._instance = super(HPUBucketingManager, cls).__new__(cls)
         return cls._instance
 
-    def initialize(self, max_num_seqs, max_num_prefill_seqs, block_size, max_num_batched_tokens, max_model_len):
+    def initialize(self,
+                   max_num_seqs,
+                   max_num_prefill_seqs,
+                   block_size,
+                   max_num_batched_tokens,
+                   max_model_len,
+                   num_speculative_tokens=0):
         self.max_num_seqs = max_num_seqs
         self.max_num_prefill_seqs = max_num_prefill_seqs
         self.block_size = block_size
         self.max_num_batched_tokens = max_num_batched_tokens
         self.num_hpu_blocks = None
         self.max_model_len = max_model_len
+        self.num_speculative_tokens = num_speculative_tokens
         self.initialized = True
         self.fallback_bs_base_step = 2
         self.fallback_seq_base_step = 32
@@ -189,6 +199,12 @@ def generate_decode_buckets(self):
                                                    self.max_num_seqs, self.max_num_prefill_seqs,
                                                    self.max_num_batched_tokens, self.block_size, self.num_hpu_blocks,
                                                    buckets_from_file)
+            if self.num_speculative_tokens:
+                # The existing buckets are used as seed decode buckets
+                self.seed_decode_buckets = self.decode_buckets
+                # More buckets are added automatically for spec decode
+                self.decode_buckets = self.generate_spec_decode_buckets(self.decode_buckets)
+
             self.log_generate_info(False)
         else:
             logger().info("Bucketing is off - skipping decode buckets generation")
@@ -232,8 +248,14 @@ def find_prompt_bucket(self, batch_size, seq_len, ctx=0):
             return found_bucket
         return (batch_size, seq_len, ctx)
 
-    def find_decode_bucket(self, batch_size, num_blocks):
+    def find_decode_bucket(self, batch_size, num_blocks, seed_buckets: bool = False):
         if self.initialized:
+            if seed_buckets and self.seed_decode_buckets is not None:
+                found_bucket = find_equal_or_closest_greater_config(self.seed_decode_buckets,
+                                                                    (batch_size, 1, num_blocks))
+                if found_bucket is not None:
+                    return found_bucket
+
             found_bucket = find_equal_or_closest_greater_config(self.decode_buckets, (batch_size, 1, num_blocks))
             if found_bucket is None:
                 new_bucket = self.generate_fallback_bucket(batch_size, 1, num_blocks)
@@ -260,6 +282,46 @@ def get_max_prompt_shape(self):
         return max(b[1] for b in self.prompt_buckets) \
                if len(self.prompt_buckets) > 0 else self.max_model_len
 
+    def generate_spec_decode_buckets(self, seed_decode_buckets):
+        max_model_len = self.max_model_len
+        block_size = self.block_size
+
+        def no_corrections(bs, query, ctx):
+            return (bs, query, ctx)
+
+        def correct_for_max_model_len(bs, query, ctx):
+            return (bs, query, min(ctx, bs * math.ceil(max_model_len / block_size)))
+
+        def get_corrector(use_contiguous_pa):
+            if use_contiguous_pa:
+                return no_corrections
+            else:
+                return correct_for_max_model_len
+
+        use_contiguous_pa = get_config().use_contiguous_pa
+        corrector = get_corrector(use_contiguous_pa)
+
+        # If spec decode enabled, generate buckets for batch_size * (1 + num_speculative_tokens)
+        num_tokens = 1 + self.num_speculative_tokens
+        buckets = set()
+        for bucket in seed_decode_buckets:
+            buckets.add(bucket)
+            bs, query, ctx = bucket
+            spec_decode_bs = bs * num_tokens
+            if spec_decode_bs <= ctx:
+                # Add a bucket with (batch_size * num_tokens, query, ctx)
+                buckets.add(corrector(spec_decode_bs, query, ctx))
+            # Add a bucket with (batch_size * num_tokens, query, ctx * num_tokens)
+            buckets.add(corrector(spec_decode_bs, query, ctx * num_tokens))
+
+        # Log the new generated spec decode buckets
+        new_buckets = sorted(buckets - set(seed_decode_buckets))
+        msg = (f"Generated {len(new_buckets)} "
+               f"spec decode buckets [bs, query, num_blocks]: {list(new_buckets)}")
+        logger().info(msg)
+
+        return sorted(buckets)
+
     @classmethod
     def get_instance(cls):
         """

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -891,11 +891,13 @@ def __init__(
                                else self.max_prefill_batch_size
         if self.enable_bucketing:
             logger.info("Bucketing is ON.")
+            num_speculative_tokens = self.speculative_config.num_speculative_tokens if self.speculative_config else 0
             self.bucketing_manager.initialize(max_num_seqs=self.max_num_seqs,
                                               max_num_prefill_seqs=max_num_prefill_seqs,
                                               block_size=self.block_size,
                                               max_num_batched_tokens=self.max_num_batched_tokens,
-                                              max_model_len=self.max_model_len)
+                                              max_model_len=self.max_model_len,
+                                              num_speculative_tokens=num_speculative_tokens)
             self.graphed_buckets: set[Any] = set()
             self.graphed_multimodal_buckets: set[Any] = set()
         else:
@@ -2031,15 +2033,18 @@ def _create_decode_input_data(self,
         # but also kvs for the current token
         num_blocks = np.ceil((context_lens + 1) / self.block_size).astype(np.int32).tolist()
 
+        num_tokens_per_req = num_scheduled_tokens[:num_decodes]
+        num_tokens = max(num_tokens_per_req)
+        # Spec decode to use seed buckets to get padded batch size
+        seek_buckets = bool(num_tokens > 1)
+
         # PAD FOR STATIC SHAPES.
         padded_batch_size: int
-        padded_batch_size = self.bucketing_manager.find_decode_bucket(num_decodes, sum(num_blocks))[0]
+        padded_batch_size = self.bucketing_manager.find_decode_bucket(num_decodes, sum(num_blocks), seek_buckets)[0]
 
         # dp aware padding
         padded_batch_size += self.get_dp_padding(padded_batch_size)
 
-        num_tokens_per_req = num_scheduled_tokens[:num_decodes]
-        num_tokens = max(num_tokens_per_req)
         total_num_scheduled_tokens = sum(num_tokens_per_req)
         num_tokens_per_req = num_tokens_per_req + [0] * (padded_batch_size - num_decodes)
 
@@ -3513,7 +3518,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
@@ -4098,7 +4105,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)}'
@@ -4177,14 +4193,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):
@@ -4324,6 +4346,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),