Add limit for decode block bucket size (#41)

vllm_hpu_extension/bucketing.py

@@ -34,6 +34,7 @@ def __init__(self, max_num_seqs, max_num_prefill_seqs, block_size,
         self.block_size = block_size
         self.max_num_batched_tokens = max_num_batched_tokens
         self._setup_buckets()
+        self.num_hpu_blocks = None
 
     def _setup_buckets(self) -> None:
         # FIXME: The default values should be max_model_len
@@ -111,8 +112,10 @@ def get_padded_prompt_seq_len(self, seq_len):
                            self.global_state.prompt_seq_bucket_cfg)
 
     def get_padded_decode_num_blocks(self, num_blocks):
-        return find_bucket(num_blocks,
-                           self.global_state.decode_block_bucket_cfg)
+        assert self.num_hpu_blocks is not None, "num_hpu_blocks is not set"
+        bucket_size = find_bucket(num_blocks,
+                                  self.global_state.decode_block_bucket_cfg)
+        return min(bucket_size, self.num_hpu_blocks)
 
     def get_padded_batch_size(self, batch_size, is_prompt):
         if is_prompt: