Buckets from file - alpha version (#375) (#496)

Cherry-pick of d84e734

---------

Signed-off-by: Agata Dobrzyniewicz <[email protected]>
Co-authored-by: Agata Dobrzyniewicz <[email protected]>
Co-authored-by: Michał Kuligowski <[email protected]>

Author: 3 people

vllm_gaudi/extension/bucketing/bucketing_file.txt

@@ -0,0 +1,24 @@
+# This is a sample bucketing file
+
+# Add buckets as follows:
+# Prompt: (batch_size, query_len, context)
+# Decode: (batch_size, 1, context)
+
+# You can also use lists to end up with Cartesian product like so:
+# (1, [256, 512], [0, 4, 8])
+# In this case you will end up with 6 buckets
+# You can also use python's range to create similiar lists
+# range(min, max, step)
+# Examples are shown below
+
+# Not supported for unified attention buckets
+# use '#' to comment out lines
+
+
+# Buckets:
+(1, 2048, 0)
+(1, [256, 512], [0, 4, 8])
+
+(64, 1, 1024)
+(32, 1, 512)
+(1, 1, range(256, 512, 32))

vllm_gaudi/extension/bucketing/common.py

@@ -54,13 +54,18 @@ def initialize(self, max_num_seqs, max_num_prefill_seqs, block_size, max_num_bat
         self.num_hpu_blocks = None
         self.max_model_len = max_model_len
         self.initialized = True
-
         self.fallback_bs_base_step = 2
         self.fallback_seq_base_step = 32
         self.fallback_blocks_base_step = 32
 
     ### GENERATE BUCKETS FUNCTIONS ###
 
+    def read_from_file(self, is_prompt):
+        file_name = get_config().VLLM_BUCKETING_FROM_FILE
+        from vllm_gaudi.extension.bucketing.file_strategy import (FileBucketingStrategy)
+        strategy = FileBucketingStrategy()
+        return strategy.get_buckets(file_name, is_prompt)
+
     def get_bucketing_strategy(self):
         strategy = None
         # TODO - we can use different strategies for decode and prompt
@@ -78,6 +83,8 @@ def get_bucketing_strategy(self):
 
     def generate_unified_buckets(self):
         if self.initialized:
+            if get_config().VLLM_BUCKETING_FROM_FILE:
+                assert "Unified attention doesn't support bucketing from file"
             from vllm_gaudi.extension.bucketing.unified import (UnifiedBucketingStrategy)
             strategy = UnifiedBucketingStrategy()
 
@@ -105,20 +112,29 @@ def generate_unified_buckets(self):
 
     def generate_prompt_buckets(self):
         if self.initialized:
-            strategy = self.get_bucketing_strategy()
-
-            bs_cfg, query_cfg, ctx_cfg = strategy.get_prompt_cfgs(max_num_prefill_seqs=self.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)
-
-            bs_range = strategy.get_range(bs_cfg)
-            query_range = strategy.get_range(query_cfg)
-            ctx_range = strategy.get_range(ctx_cfg)
+            buckets_from_file = None
+            bs_range = []
+            query_range = []
+            ctx_range = []
+            if get_config().VLLM_BUCKETING_FROM_FILE:
+                buckets_from_file = self.read_from_file(is_prompt=True)
+            else:
+                strategy = self.get_bucketing_strategy()
+
+                bs_cfg, query_cfg, ctx_cfg = strategy.get_prompt_cfgs(
+                    max_num_prefill_seqs=self.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)
+
+                bs_range = strategy.get_range(bs_cfg)
+                query_range = strategy.get_range(query_cfg)
+                ctx_range = strategy.get_range(ctx_cfg)
 
             self.prompt_buckets = generate_buckets(bs_range, query_range, ctx_range, True, self.max_model_len,
                                                    self.max_num_seqs, self.max_num_prefill_seqs,
-                                                   self.max_num_batched_tokens, self.block_size, self.num_hpu_blocks)
+                                                   self.max_num_batched_tokens, self.block_size, self.num_hpu_blocks,
+                                                   buckets_from_file)
             self.log_generate_info(True)
         else:
             logger().info("Bucketing is off - skipping prompt buckets generation")
@@ -127,24 +143,33 @@ def generate_prompt_buckets(self):
 
     def generate_decode_buckets(self):
         if self.initialized:
-            strategy = self.get_bucketing_strategy()
-
-            bs_cfg, query_cfg, ctx_cfg = strategy.get_decode_cfgs(max_num_seqs=self.max_num_seqs,
-                                                                  block_size=self.block_size,
-                                                                  max_num_batched_tokens=self.max_num_batched_tokens,
-                                                                  max_model_len=self.max_model_len,
-                                                                  max_blocks=self.num_hpu_blocks)
-
-            bs_range = strategy.get_range(bs_cfg)
-            query_range = strategy.get_range(query_cfg)
-            ctx_range = strategy.get_range(ctx_cfg)
-
-            if get_config().use_contiguous_pa and ctx_range[-1] < self.num_hpu_blocks:
-                ctx_range.append(self.num_hpu_blocks)
+            buckets_from_file = None
+            bs_range = []
+            query_range = []
+            ctx_range = []
+            if get_config().VLLM_BUCKETING_FROM_FILE:
+                buckets_from_file = self.read_from_file(is_prompt=False)
+            else:
+                strategy = self.get_bucketing_strategy()
+
+                bs_cfg, query_cfg, ctx_cfg = strategy.get_decode_cfgs(
+                    max_num_seqs=self.max_num_seqs,
+                    block_size=self.block_size,
+                    max_num_batched_tokens=self.max_num_batched_tokens,
+                    max_model_len=self.max_model_len,
+                    max_blocks=self.num_hpu_blocks)
+
+                bs_range = strategy.get_range(bs_cfg)
+                query_range = strategy.get_range(query_cfg)
+                ctx_range = strategy.get_range(ctx_cfg)
+
+                if get_config().use_contiguous_pa and ctx_range[-1] < self.num_hpu_blocks:
+                    ctx_range.append(self.num_hpu_blocks)
 
             self.decode_buckets = generate_buckets(bs_range, query_range, ctx_range, False, self.max_model_len,
                                                    self.max_num_seqs, self.max_num_prefill_seqs,
-                                                   self.max_num_batched_tokens, self.block_size, self.num_hpu_blocks)
+                                                   self.max_num_batched_tokens, self.block_size, self.num_hpu_blocks,
+                                                   buckets_from_file)
             self.log_generate_info(False)
         else:
             logger().info("Bucketing is off - skipping decode buckets generation")
@@ -225,8 +250,17 @@ def get_bucketing_manager():
     return instance
 
 
-def generate_buckets(bs_range, query_range, ctx_range, is_prompt, max_model_len, max_num_seqs, max_num_prefill_seqs,
-                     max_num_batched_tokens, block_size, max_blocks):
+def generate_buckets(bs_range,
+                     query_range,
+                     ctx_range,
+                     is_prompt,
+                     max_model_len,
+                     max_num_seqs,
+                     max_num_prefill_seqs,
+                     max_num_batched_tokens,
+                     block_size,
+                     max_blocks,
+                     file_buckets=None):
     use_merged_prefill = get_config().merged_prefill
     use_contiguous_pa = get_config().use_contiguous_pa
 
@@ -307,15 +341,23 @@ def get_filters(is_prompt, use_merged_prefill, use_contiguous_pa):
     buckets_2d = set()
     omitted_buckets = set()
     filters = get_filters(is_prompt, use_merged_prefill, use_contiguous_pa)
-    for bs_idx, bs in enumerate(bs_range):
-        for query_idx, query in enumerate(query_range):
-            buckets_2d.update(
-                expand_to_neighbor_buckets(bs_idx, bs_range, query_idx, query_range, max_num_batched_tokens))
-
-    for bs, query in buckets_2d:
-        for ctx in ctx_range:
-            if all(bucket_filter(bs, query, ctx) for bucket_filter in filters):
-                buckets.add((bs, query, ctx))
+
+    if file_buckets:
+        for bs, query, blocks in file_buckets:
+            if all(bucket_filter(bs, query, blocks) for bucket_filter in filters):
+                buckets.add((bs, query, blocks))
+    else:
+        for bs_idx, bs in enumerate(bs_range):
+            for ctx_idx, ctx in enumerate(ctx_range):
+                local_buckets = expand_to_neighbor_buckets(bs_idx, bs_range, ctx_idx, ctx_range,
+                                                           max_num_batched_tokens) if not is_prompt else {(bs, ctx)}
+                buckets_2d.update(local_buckets)
+
+        for bs, ctx in buckets_2d:
+            for query in query_range:
+                if all(bucket_filter(bs, query, ctx) for bucket_filter in filters):
+                    buckets.add((bs, query, ctx))
+
     if not buckets:
         phase = 'prompt' if is_prompt else 'decode'
         for bucket in omitted_buckets:

vllm_gaudi/extension/bucketing/file_strategy.py

@@ -0,0 +1,54 @@
+import itertools
+import operator
+import os
+import math
+import ast
+from dataclasses import dataclass, field
+from typing import List, Tuple
+
+from vllm_gaudi.extension.logger import logger as logger
+from vllm_gaudi.extension.runtime import get_config
+
+
+class FileBucketingStrategy:
+
+    def get_buckets(self, file_name, is_prompt):
+        prompt_buckets = []
+        decode_buckets = []
+
+        with open(file_name, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith('#'):
+                    continue
+
+                try:
+                    bucket = eval(line, {"__builtins__": None}, {"range": range})
+                except Exception as e:
+                    print(f"Skipping line due to eval error: {e} - {line}")
+                    continue
+
+                if not isinstance(bucket, tuple) or len(bucket) != 3:
+                    print('Skipping line due to incorrect format - ', bucket)
+                    continue
+
+                x_num = ensure_is_list(bucket[0])
+                y_num = ensure_is_list(bucket[1])
+                z_num = ensure_is_list(bucket[2])
+
+                for full_bucket in itertools.product(x_num, y_num, z_num):
+                    x, y, z = map(int, full_bucket)
+                    if y == 1:
+                        decode_buckets.append((x, y, z))
+                    else:
+                        prompt_buckets.append((x, y, z))
+        return sorted(prompt_buckets) if is_prompt else sorted(decode_buckets)
+
+
+def ensure_is_list(value):
+    if isinstance(value, list):
+        return value
+    elif isinstance(value, range):
+        return list(value)
+    else:
+        return [value]

vllm_gaudi/extension/features.py

@@ -34,6 +34,7 @@ def get_user_flags():
         Env('VLLM_DECODE_BLOCK_BUCKET_STEP', int),
         Env('VLLM_DECODE_BLOCK_BUCKET_MAX', int),
         Env('VLLM_DECODE_BLOCK_BUCKET_LIMIT', int),
+        Env('VLLM_BUCKETING_FROM_FILE', str),
 
         # Non-vllm flags that are also important to print
         Env('EXPERIMENTAL_WEIGHT_SHARING', str),