Add option to limit number of buckets

vllm/worker/habana_model_runner.py

@@ -15,6 +15,7 @@
                     Optional, Set, Tuple, Type, TypeVar, Union)
 
 import habana_frameworks.torch as htorch
+import numpy as np
 import torch
 
 from vllm.attention import AttentionMetadata, get_attn_backend
@@ -60,7 +61,7 @@ def read_bucket_settings(phase: str, dim: str, **defaults):
     param is either 'min', 'step' or 'max'
     example env variable: VLLM_DECODE_BS_BUCKET_STEP=128
     """
-    params = ['min', 'step', 'max']
+    params = ['min', 'step', 'max', 'limit']
     values = [
         int(
             os.environ.get(f'VLLM_{phase}_{dim}_BUCKET_{p}'.upper(),
@@ -94,14 +95,169 @@ def warmup_range(config: Tuple[int, int, int]):
     return list(ramp_up_tw) + list(stable)
 
 
-def warmup_buckets(bs_bucket_config, seq_bucket_config,
-                   max_num_batched_tokens):
-    buckets = itertools.product(warmup_range(bs_bucket_config),
-                                warmup_range(seq_bucket_config))
-    # Remove buckets exceeding batch token budget
-    filtered_buckets = filter(
-        lambda bucket: bucket[0] * bucket[1] <= max_num_batched_tokens,
-        buckets)
+def warmup_range_with_limit(config: Tuple[int, int, int, int], fill=True):
+    """ 
+    NOTE(kzawora): we'll use exponential spacing for buckets in which scaled 
+    power will return bmin for first bucket iteration, and bmax for last 
+    iteration, with elements between determined by the exponent, and base being 
+    unchanged. Note that after padding to bstep, duplicates may occur. 
+    Handling of duplicates is configured by fill parameter. 
+
+    If fill is False, duplicates are removed and less buckets are returned. 
+
+    If fill is True, duplicates are resolved by selecting the closest (larger 
+    or smaller) bucket. If duplicate resolution is not possible, less buckets 
+    are returned. In that case, buckets are guaranteed to be linearly spaced.
+
+    Example (bmin=128, bstep=128, bmax=2048, num_buckets=10):
+
+    There are 16 possible buckets (2048/128), and we'll attempt to select 10 of 
+    them with exponential spacing.
+    base = (bmax/bmin) ** (1/(num_buckets-1)); (2048/128) ** (1/9) = 1.36079
+    exponent = i
+    power = base ** exponent
+    scaled_power = b_min * power
+
+    For i == 0 (first bucket), power is 1.36079 ** 0 = 1; 
+        scaled_power is 1 * 128 = 128 (==bmin)
+    For i == 9 (last bucket), power is 1.36079 ** 9 = 16; 
+        scaled_power is 16 * 128 = 2048 (==bmax)
+
+    So, computing for all buckets:
+    scaled_powers_unpadded     = [bmin*base^0(==bmin), bmin*base^1, bmin*base^2,       ...,     bmin*base^9(==bmax)]
+    scaled_powers_unpadded     = [128.00, 174.18, 237.02, 322.54, 438.91, 597.26, 812.75, 1105.98, 1505.01, 2048.00]
+
+    if fill is False:
+        scaled_powers_padded   = [   128,    256,    256,    384,    512,    640,    896,    1152,    1536,    2048]
+                                               ^_______^ 
+                                               duplicates
+
+        buckets                = [   128,    256,            384,    512,    640,    896,    1152,    1536,    2048]
+                                                      ^ 
+                                         duplicate bucket removed
+
+        len(buckets) = 9, num_buckets = 10
+
+    if fill is True:
+        buckets                = [   128,    256,    384,    512,    640,    768,    896,    1152,    1536,    2048]
+                                                      ^_______^_______^_______^ 
+                                                   closest unused buckets selected
+                                                              ^_______^_______^ 
+                                      these become duplicates once previous duplicates are resolved
+
+        In this case we'll have four duplicated buckets:
+
+        174.18 -> 256, optimal bucket,
+        237.02 -> (256) -> 384, taking closest available bucket, 
+            as optimal bucket 256 was already captured by 174.18, 
+        322.54 -> (384) -> 512, taking closest available bucket, 
+            as optimal bucket 384 was already captured by 237.02,
+        438.91 -> (512) -> 640, taking closest available bucket, 
+            as optimal bucket 512 was already captured by 322.54,
+        597.26 -> (640) -> 768, taking closest available bucket, 
+            as optimal bucket 640 was already captured by 438.91,
+        812.75 -> 896, optimal bucket
+
+        len(buckets) = 10, num_buckets = 10
+
+        In this case, the end result has the same buckets as fill=False, 
+        but with additional bucket 768 added. 
+        The difference is more pronounced for larger ranges and larger number 
+        of buckets.
+
+    """ # noqa: E501
+
+    bmin, bstep, bmax, num_buckets = config
+    linear_buckets = set(np.arange(bmin, bmax + 1, step=bstep))
+    assert num_buckets > 0, "num_buckets must be a positive integer"
+    if num_buckets == 1:
+        return [bmax]
+    buckets: Set[Tuple[int, int]] = set()
+    for i in range(num_buckets):
+        power_unpadded = bmin * np.float_power(
+            bmax / bmin, (1. / float(num_buckets - 1)) * i)
+        bucket = math.ceil(power_unpadded / bstep) * bstep
+        if fill and bucket in buckets:
+            available_buckets = linear_buckets.difference(buckets)
+            if len(available_buckets) == 0:
+                break  # there are no more unique buckets, let's exit now
+            new_bucket = min(available_buckets,
+                             key=lambda x: abs(x - power_unpadded))
+            buckets.add(new_bucket)
+        else:
+            buckets.add(bucket)
+    return list(sorted(buckets))
+
+
+def warmup_buckets(bs_bucket_config, seq_bucket_config, max_num_batched_tokens,
+                   block_size, max_model_len):
+    # NOTE(kzawora): If either max_model_len or max_num_batched_tokens are
+    # not divisible by block_size, we round them up here to block_size
+
+    bs_buckets = warmup_range(bs_bucket_config[:3])
+    seq_len_buckets = warmup_range(seq_bucket_config[:3])
+    if bs_bucket_config[3] != 0 and len(bs_buckets) > bs_bucket_config[3]:
+        bs_buckets = warmup_range_with_limit(bs_bucket_config)
+    if seq_bucket_config[3] != 0 and len(
+            seq_len_buckets) > seq_bucket_config[3]:
+        seq_len_buckets = warmup_range_with_limit(seq_bucket_config)
+
+    buckets = list(itertools.product(bs_buckets, seq_len_buckets))
+
+    # Remove buckets that either exceed token budget
+    # or are greater than max model length
+    filtered_buckets = list(
+        filter(
+            lambda bucket: bucket[0] * bucket[1] <= math.ceil(
+                max_num_batched_tokens / block_size) * block_size and bucket[1]
+            <= max_model_len, buckets))
+
+    # NOTE(kzawora): We need to make sure we cover the largest valid scenario
+    # (max_num_batched_tokens/bs). If we drop out-of-bound buckets, we might
+    # be left with last bucket that's too small and encounter recompilations
+    # when exceeding that bucket and still being in the valid range.
+
+    # Example:
+    # block_size = 128
+    # max_model_len = 768 (6 blocks)
+    # max_num_batched_tokens = 1536 (12 blocks)
+    # bs_buckets = [1, 2, 4, 5]
+    # seq_len_buckets = [128, 256, 512, 1024]
+    # buckets = [
+    #   (1,128), (1,256), (1,512), (1,1024), # last one is invalid, 1024 > 768
+    #   (2,128), (2,256), (2,512), (2,1024)  # last one is invalid, 2048 > 1536
+    #   (4,128), (4,256), (4,512), (4,1024)  # last two are invalid, 2048 > 1536
+    #   (5,128), (5,256), (5,512), (5,1024)  # last two are invalid, 2560 > 1536
+    # ]
+    # filtered_buckets = [
+    #   (1,128), (1,256), (1,512), # 512-768 is valid and not bucketed
+    #   (2,128), (2,256), (2,512), # 512-768 is valid and not bucketed
+    #   (4,128), (4,256),          # 256-384 is valid and not bucketed
+    #   (5,128), (5,256),          # 256-307 is valid and not bucketed
+    # ]
+    # boundary_buckets = [
+    #   (1,768), # covers 512-768 range, constrained by max_model_len
+    #   (2,768)  # covers 512-768 range, constrained by budget and max_model_len
+    #   (4,384)  # covers 256-384 range, constrained by token budget
+    #   (5,384)  # covers 256-307 range, constrained by token budget,
+    #            #                       rounded up to block_size(=128)
+    # ]
+    # While (5,384) bucket exceeds the token budget of 1536 (5*386=1920),
+    # it is the smallest bucket that can handle all valid use cases on bs=5,
+    # as all buckets for decode phase need to be aligned to block_size.
+
+    # Generate the largest valid bucket for all valid batch sizes.
+    # DO NOT use bs_buckets in here as it can include invalid batches
+    valid_bs_buckets = set(list(zip(*filtered_buckets))[0])
+    boundary_buckets = [
+        (bs,
+         min(
+             math.ceil(max_model_len / block_size) * block_size,
+             math.ceil(max_num_batched_tokens / (bs * block_size)) *
+             block_size)) for bs in valid_bs_buckets
+    ]
+    filtered_buckets = list(set([*filtered_buckets, *boundary_buckets]))
+
     return list(
         sorted(filtered_buckets, key=lambda b: (b[0] * b[1], b[1], b[0])))
 
@@ -118,15 +274,6 @@ def round_up(value: int, k: int):
     return (value + k - 1) // k * k
 
 
-def find_bucket(value: int, config: Tuple[int, int, int]):
-    bmin, bstep, bmax = config
-    if value < bstep:
-        result = min(next_pow2(value), bstep)
-    else:
-        result = round_up(value, bstep)
-    return result
-
-
 def subtuple(obj: object,
              typename: str,
              to_copy: List[str],
@@ -509,31 +656,39 @@ def _setup_buckets(self) -> None:
                                                          step=32,
                                                          max=min(
                                                              self.max_num_seqs,
-                                                             max_bucket_cfg))
+                                                             max_bucket_cfg),
+                                                         limit=0)
         self.decode_bs_bucket_cfg = read_bucket_settings('decode',
                                                          'bs',
                                                          min=1,
                                                          step=128,
-                                                         max=self.max_num_seqs)
-        self.prompt_seq_bucket_cfg = read_bucket_settings('prompt',
-                                                          'seq',
-                                                          min=self.block_size,
-                                                          step=self.block_size,
-                                                          max=1024)
-        self.decode_seq_bucket_cfg = read_bucket_settings('decode',
-                                                          'seq',
-                                                          min=self.block_size,
-                                                          step=self.block_size,
-                                                          max=2048)
+                                                         max=self.max_num_seqs,
+                                                         limit=0)
+        self.prompt_seq_bucket_cfg = read_bucket_settings(
+            'prompt',
+            'seq',
+            min=self.block_size,
+            step=self.block_size,
+            max=self.max_model_len,
+            limit=0)
+        self.decode_seq_bucket_cfg = read_bucket_settings(
+            'decode',
+            'seq',
+            min=self.block_size,
+            step=self.block_size,
+            max=self.max_model_len,
+            limit=0)
         self.graphed_buckets: Set[Any] = set()
 
-        msg = ("Prompt bucket config (min, step, max_warmup) "
+        msg = ("Prompt bucket config (min, step, max_warmup, limit) "
                f"bs:{self.prompt_bs_bucket_cfg}, "
                f"seq:{self.prompt_seq_bucket_cfg}")
         logger.info(msg)
         self.prompt_buckets = warmup_buckets(self.prompt_bs_bucket_cfg,
                                              self.prompt_seq_bucket_cfg,
-                                             self.max_num_batched_tokens)
+                                             self.max_num_batched_tokens,
+                                             self.block_size,
+                                             self.max_model_len)
 
         if self.lora_config:
             self.prompt_buckets[:] = [
@@ -545,18 +700,32 @@ def _setup_buckets(self) -> None:
                f"prompt buckets: {list(sorted(self.prompt_buckets))}")
         logger.info(msg)
 
-        msg = ("Decode bucket config (min, step, max_warmup) "
+        msg = ("Decode bucket config (min, step, max_warmup, limit) "
                f"bs:{self.decode_bs_bucket_cfg}, "
                f"seq:{self.decode_seq_bucket_cfg}")
         logger.info(msg)
         self.decode_buckets = warmup_buckets(self.decode_bs_bucket_cfg,
                                              self.decode_seq_bucket_cfg,
-                                             self.max_num_batched_tokens)
+                                             self.max_num_batched_tokens,
+                                             self.block_size,
+                                             self.max_model_len)
         if self.lora_config:
             self.decode_buckets[:] = [
                 bucket for bucket in self.decode_buckets
                 if self._is_valid_bucket(bucket)
             ]
+
+        find_bucket = lambda bucket, x: next(p for p in sorted(bucket)
+                                             if p >= x)
+        get_bucket_dim = lambda bucket, dim: [b[dim] for b in bucket]
+        self.find_bs_bucket = lambda x, is_prompt: find_bucket(
+            get_bucket_dim(
+                self.prompt_buckets
+                if is_prompt else self.decode_buckets, 0), x)
+        self.find_seq_bucket = lambda x, is_prompt: find_bucket(
+            get_bucket_dim(
+                self.prompt_buckets
+                if is_prompt else self.decode_buckets, 1), x)
         msg = (f"Generated {len(self.decode_buckets)} decode buckets: "
                f"{list(sorted(self.decode_buckets))}")
         logger.info(msg)
@@ -688,9 +857,8 @@ def _prepare_prompt(
             multi_modal_input = None
 
         max_prompt_block_table_len = max(len(t) for t in prefix_block_tables)
-        max_prompt_len = max(
-            find_bucket(max(seq_lens), self.prompt_seq_bucket_cfg),
-            self.block_size)
+        max_prompt_len = max(self.find_seq_bucket(max(seq_lens), True),
+                             self.block_size)
 
         for seq_group_metadata, context_len in zip(seq_group_metadata_list,
                                                    context_lens):
@@ -896,9 +1064,7 @@ def prepare_input_tensors(
         self.profiler.start('internal', base_event_name)
 
         real_batch_size = len(seq_group_metadata_list)
-        bucket_cfg = self.prompt_bs_bucket_cfg if is_prompt else \
-            self.decode_bs_bucket_cfg
-        batch_size_padded = find_bucket(real_batch_size, bucket_cfg)
+        batch_size_padded = self.find_bs_bucket(real_batch_size, is_prompt)
         batch_size_padding = batch_size_padded - real_batch_size
         seq_group_metadata_list = seq_group_metadata_list.copy()
         seq_group_metadata_list.extend(seq_group_metadata_list[0]
@@ -1087,16 +1253,30 @@ def create_dummy_seq_group_metadata(self,
     def profile_run(self) -> None:
         num_layers = self.model_config.get_num_layers(self.parallel_config)
         kv_caches = [None] * num_layers
-        max_batch_size = self.prompt_bs_bucket_cfg[-1]
-        max_seq_len = self.prompt_seq_bucket_cfg[-1]
-        if self.lora_config:
-            max_seq_len = self.max_num_batched_tokens // max_batch_size
 
-        self.warmup_scenario(max_batch_size,
-                             max_seq_len,
-                             True,
-                             kv_caches,
-                             is_profile_run=True)
+        # take two largest buckets in terms of token counts
+        # first with largest batch size, second with largest seq length
+        # if both are the same, only single warmup run is executed
+        warmup_scenarios = set()
+        warmup_scenarios.add(
+            max(self.prompt_buckets,
+                key=lambda item: (item[0], item[0] * item[1])))
+        warmup_scenarios.add(
+            max(self.prompt_buckets,
+                key=lambda item: (item[1], item[0] * item[1])))
+        for idx, (max_batch_size, max_seq_len) in enumerate(warmup_scenarios):
+            msg = (f'[{idx+1}/{len(warmup_scenarios)}] '
+                   f'Executing profile run for bs={max_batch_size}, '
+                   f'seq_len={max_seq_len}, '
+                   f'num_batched_tokens={max_batch_size*max_seq_len}')
+            logger.info(msg)
+            if self.lora_config:
+                max_seq_len = self.max_num_batched_tokens // max_batch_size
+            self.warmup_scenario(max_batch_size,
+                                 max_seq_len,
+                                 True,
+                                 kv_caches,
+                                 is_profile_run=True)
 
     def warmup_scenario(self,
                         batch_size,