Limit decode block size

vllm/worker/hpu_model_runner.py

@@ -227,10 +227,11 @@ def round_up(value: int, k: int):
 
 
 def find_bucket(value: int, config: Tuple[int, int, int]):
-    bmin, bstep, _ = config
+    bmin, bstep, bmax = config
     next_step = round_up(value, bstep)
     next_pow = next_pow2(value, bmin)
-    return max(bmin, min(next_step, next_pow))
+    b = max(bmin, min(next_step, next_pow))
+    return min(b, bmax)
 
 
 def align_workers(value, op):
@@ -808,8 +809,11 @@ def _setup_buckets(self) -> None:
             'block',
             min=self.block_size,
             step=self.block_size,
-            max=max(self.block_size,
-                    self.max_num_seqs * max_decode_seq // self.block_size))
+            max=min(max(self.block_size,
+                        self.max_num_seqs * max_decode_seq // self.block_size),
+                    self.cache_config.num_gpu_blocks)
+            )
+
         self.graphed_buckets: Set[Any] = set()
 
         msg = ("Prompt bucket config (min, step, max_warmup) "