bucket: add query len 1 to prefill bucket

Signed-off-by: Xinyu Chen <[email protected]>

Author: xinyu-intel

tests/unit_tests/test_bucketing.py

@@ -39,6 +39,12 @@ def test_warmup_range():
     assert result == [2, 4, 8, 16, 32, 64, 128]
 
 
+def test_warmup_range_with_one():
+    config = (1, 64, 128)
+    result = linear.warmup_range(config)
+    assert result == [1, 2, 4, 8, 16, 32, 64, 128]
+
+
 def test_generate_prompt_buckets():
     max_num_batched_tokens = 2048
     block_size = 64

vllm_gaudi/extension/bucketing/exponential.py

@@ -36,7 +36,7 @@ def get_prompt_cfgs(self, max_num_prefill_seqs, block_size, max_num_batched_toke
         prompt_bs_limit = math.ceil(math.log2(max_num_prefill_seqs)) + 1
         prompt_bs_bucket_cfg = [1, 2, max_num_prefill_seqs, prompt_bs_limit]
         max_prompt_seq_limit = math.ceil(math.log2(max_num_batched_tokens))
-        prompt_query_bucket_cfg = [block_size, block_size, max_num_batched_tokens, max_prompt_seq_limit]
+        prompt_query_bucket_cfg = [1, block_size, max_num_batched_tokens, max_prompt_seq_limit]
         max_ctx = max(1, math.ceil((max_model_len - prompt_query_bucket_cfg[0]) // block_size))
         max_prompt_ctx_limit = 2 if max_ctx == 1 else math.ceil(math.log2(max_ctx)) + 1
         prompt_ctx_bucket_cfg = [0, 1, max_ctx, max_prompt_ctx_limit]
@@ -125,7 +125,8 @@ def warmup_range_with_limit(config: Tuple[int, int, int, int], long_context=Fals
 
     bmin, bstep, bmax, num_buckets = config
     add_zero_bucket = bmin == 0
-    if add_zero_bucket:
+    add_one_bucket = bmin == 1
+    if add_zero_bucket or add_one_bucket:
         bmin = bstep
     linear_buckets = set(np.arange(bmin, bmax + 1, step=bstep))
     assert num_buckets > 0, "num_buckets must be a positive integer"
@@ -176,4 +177,6 @@ def warmup_range_with_limit(config: Tuple[int, int, int, int], long_context=Fals
                 buckets.add(bucket)
     if add_zero_bucket:
         buckets.add(0)
+    if add_one_bucket:
+        buckets.add(1)
     return list(sorted(buckets))

vllm_gaudi/extension/bucketing/linear.py

@@ -18,7 +18,7 @@ def get_prompt_cfgs(self, max_num_prefill_seqs, block_size, max_num_batched_toke
         prompt_bs_bucket_cfg = read_bucket_settings('prompt', 'bs', min=1, step=1, max=max_num_prefill_seqs)
         prompt_query_bucket_cfg = read_bucket_settings('prompt',
                                                        'query',
-                                                       min=block_size,
+                                                       min=1,
                                                        step=block_size,
                                                        max=max_num_batched_tokens)
         max_ctx = math.ceil((max_model_len - prompt_query_bucket_cfg[0]) // block_size)

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -1974,10 +1974,10 @@ def _form_unified_prefill_batch(self, contents):
 
     def _create_dummy_prefill_batch_contents(self, num_prefills: int) -> list[PrefillInputData]:
         req_id = str(-1)
-        context_len = 0
-        query_len = 128
+        context_len = 127 if has_kv_transfer_group() else 0
+        query_len = 1 if has_kv_transfer_group() else 128
         prompt_tokens = 128
-        token_ids = list(int(i) for i in range(prompt_tokens))
+        token_ids = list(int(i) for i in range(query_len))
         num_blocks = round_up(context_len + query_len, self.block_size) // self.block_size
         blocks = [0] * num_blocks
         num_output_logits = context_len + query_len - prompt_tokens + 1