Enable modification of prompt BS (#258)

Enable modification of prefill BS with usage of
`VLLM_PROMPT_BS_BUCKET_MAX` environment variable. The default size of
prefill BS is set to 1 (remains the same as it was before the change).
cherry-pick: #224

---------

Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>

Author: ksmusz

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -23,7 +23,7 @@
                                            format_bytes, setup_profiler)
 from vllm_gaudi.extension.runtime import finalize_config, get_config
 from vllm_gaudi.extension.unified import (create_unified_batch)
-from vllm_gaudi.extension.utils import pad_list
+from vllm_gaudi.extension.utils import pad_list, with_default
 from vllm_gaudi.extension.debug import init_debug_logger
 
 from vllm.attention.backends.abstract import AttentionType
@@ -807,8 +807,7 @@ def __init__(
         self.use_hpu_graph = not self.model_config.enforce_eager
         self.max_batch_size = self.scheduler_config.max_num_seqs
         self.max_num_seqs = self.scheduler_config.max_num_seqs
-        # TODO(kzawora): add knob for that
-        self.max_prefill_batch_size = 1
+        self.max_prefill_batch_size = with_default(get_config().VLLM_PROMPT_BS_BUCKET_MAX, 1)
         self.seen_configs: set = set()
         self.max_num_batched_tokens = \
             self.scheduler_config.max_num_batched_tokens
@@ -1684,6 +1683,10 @@ def _form_prefill_batch(self, contents):
         # for the valid tokens before padding.
         # This would require getting multimodal input embeddings here as well
         token_ids = self._align_and_pad(contents.token_ids, (target_bs, target_seq), itertools.repeat(-1))
+        # Update query_lens and context_lens after padding
+        query_lens.extend([0] * (target_bs - len(query_lens)))
+        context_lens.extend([0] * (target_bs - len(context_lens)))
+
         # If the model uses M-RoPE, we need to fill
         # and pad the M-RoPE positions for the scheduled prefill tokens
         if self.uses_mrope:
@@ -3819,13 +3822,14 @@ def _prepare_dummy_scenario(self, prompt_cfg, decode_cfg):
                 prompt_total_tokens, prompt_num_context_blocks = \
                      self.get_merged_prefill_seq_lens(prompt_query_len,
                                                  prompt_num_blocks)
-            for tokens, context_len in zip(prompt_total_tokens, prompt_num_context_blocks):
-                self._add_dummy_request(requests,
-                                        scheduled_tokens,
-                                        num_computed_tokens=(context_len * self.block_size),
-                                        total_tokens=tokens,
-                                        scheduled_tokens=prompt_query_len,
-                                        is_prompt=True)
+            for _ in range(prompt_bs):
+                for tokens, context_len in zip(prompt_total_tokens, prompt_num_context_blocks):
+                    self._add_dummy_request(requests,
+                                            scheduled_tokens,
+                                            num_computed_tokens=(context_len * self.block_size),
+                                            total_tokens=tokens,
+                                            scheduled_tokens=prompt_query_len,
+                                            is_prompt=True)
         if decode_cfg:
             decode_bs, decode_query_len, decode_num_blocks = decode_cfg
             if self.use_contiguous_pa: