Unified attention improvemets (#363)

- [x] warmup funcioning
- [x] no recompiles

---------

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

Author: 3 people

vllm_gaudi/extension/bucketing/common.py

@@ -340,12 +340,12 @@ def generate_unified_buckets(query_range, shared_ctx_range, unique_ctx_range, bs
             max_bs = min(bs, query)
             if math.ceil(shared_ctx * block_size // max_bs) <= max_model_len:
                 buckets.add((query, shared_ctx, unique_ctx, causal))
-        elif (query <= bs):
+        elif query <= bs:
             # non causal query = current bs
             if shared_ctx > 0 or unique_ctx > 0:
-                if shared_ctx == 0 or (query > 1 and \
-                    math.ceil(shared_ctx * block_size // (query // 2)) <= max_model_len):
-                    buckets.add((query, shared_ctx, unique_ctx, causal))
+                if shared_ctx == 0 or (math.ceil(shared_ctx * block_size // (query // 2)) <= max_model_len):
+                    if shared_ctx > 0 or query <= unique_ctx:
+                        buckets.add((query, shared_ctx, unique_ctx, causal))
 
     return sorted(buckets)
 

vllm_gaudi/extension/bucketing/unified.py

@@ -14,8 +14,8 @@ class UnifiedBucketingStrategy():
 
     def get_unified_cfgs(self, bs, max_model_len, block_size, max_blocks, max_num_batched_tokens):
         # [min, max, turning_point]
-        query_cfg = [block_size, max_num_batched_tokens, bs]
-        max_shared_ctx = math.ceil(max_model_len // block_size) * bs
+        query_cfg = [1, max_num_batched_tokens, bs]
+        max_shared_ctx = min(math.ceil(max_model_len // block_size), max_blocks)
         shared_ctx_cfg = [0, max_shared_ctx, bs]
         max_unique_ctx = max_blocks
         unique_ctx_cfg = [0, max_unique_ctx, bs]
@@ -28,19 +28,24 @@ def get_range(self, cfg):
 
 def warmup_unified_range(cfg):
     bmin, bmax, turning_point = cfg
+    limit = 10
+    round_up = 128
 
     buckets: Set[Tuple[int, int]] = set()
 
     if bmin == 0:
         buckets.add(bmin)
-
-    # alpha version: [bs/4, bs/2, bs, bt/4, bt/2, bt]
-
-    buckets.add(turning_point // 4)
-    buckets.add(turning_point // 2)
-    buckets.add(turning_point)
-    buckets.add(bmax // 4)
-    buckets.add(bmax // 2)
-    buckets.add(bmax)
+        bmin = 1
+
+    num_buckets_exp = limit
+    first_step = bmax
+
+    for i in range(num_buckets_exp):
+        power_unpadded = bmin * np.float_power(first_step / bmin, (1. / float(num_buckets_exp - 1)) * i)
+        if i == limit - 1:
+            bucket = bmax
+        else:
+            bucket = math.ceil(power_unpadded / round_up) * round_up
+        buckets.add(bucket)
 
     return list(sorted(buckets))

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -863,16 +863,8 @@ def unified_bucketing_fn(self, is_causal, query_len, shared_blocks, unique_block
         if not get_config().use_bucketing:
             return query_len, shared_blocks, unique_blocks, logits
 
-        def bucketize(x, buckets):
-            if x < buckets[-1]:
-                return next(b for b in buckets if b >= x)
-            else:
-                return round_up(x, buckets[-1])
-
-        logits_buckets = [self.max_num_seqs]
-        logits = min(bucketize(logits, logits_buckets), query_len)
         new_bucket = self.bucketing_manager.find_unified_bucket(query_len, shared_blocks, unique_blocks, is_causal)
-        return (new_bucket[0], new_bucket[1], new_bucket[2], logits)
+        return (new_bucket[0], new_bucket[1], new_bucket[2], self.max_num_seqs)
 
     def create_lora_mask(self, input_tokens: torch.Tensor, lora_ids: list[int], is_prompt: bool):
         '''
@@ -1491,7 +1483,7 @@ def _generate_req_id_output_token_ids_lst(self,
             # Merged prefill case: remove requests without logits
             req_id_output_token_ids_lst = [r for r in req_id_output_token_ids_lst if r[0] in logits_reqs]
         else:
-            if pad_to is not None:
+            if pad_to is not None and len(req_id_output_token_ids_lst) > 0:
                 while len(req_id_output_token_ids_lst) < pad_to:
                     req_id_output_token_ids_lst.append(req_id_output_token_ids_lst[0])
         return req_id_output_token_ids_lst
@@ -3858,12 +3850,10 @@ def _prepare_dummy_unified_scenario(self, unified_cfg):
             for query, blocks in zip(prompt_reqs_query, prompt_reqs_blocks):
                 self._add_dummy_unified_request(requests, True, False, blocks, num_computed_tokens, query,
                                                 scheduled_tokens)
-
         else:
             remaining_samples = query_len
             base = shared_ctx_len // remaining_samples
             remain = shared_ctx_len % remaining_samples
-
             all_shared_blocks_ids = [block for block in range(shared_ctx_len)]
             unique_block = unique_ctx_len - 1
             # do not use unique block id
@@ -3887,8 +3877,16 @@ def _prepare_dummy_unified_scenario(self, unified_cfg):
                     split_shared_blocks_ids[target].append(block)
 
             # add unique id
-            min_idx = min(range(remaining_samples), key=lambda j: len(split_shared_blocks_ids[j]))
-            split_shared_blocks_ids[min_idx].append(unique_block)
+            if unique_ctx_len > 0:
+                min_idx = min(range(remaining_samples), key=lambda j: len(split_shared_blocks_ids[j]))
+                split_shared_blocks_ids[min_idx].append(unique_block)
+
+            for i in range(len(split_shared_blocks_ids)):
+                if not split_shared_blocks_ids[i]:
+                    if unique_block - i >= 0:
+                        split_shared_blocks_ids[i] = [unique_block - i]
+                    else:
+                        split_shared_blocks_ids[i] = [all_shared_blocks_ids[0]]
 
             for request_blocks in split_shared_blocks_ids:
                 self._add_dummy_unified_request(requests, False, False, request_blocks, num_computed_tokens, 1,