[UT] Fix test args for bucketing tests (#105)

Signed-off-by: Agata Dobrzyniewicz <[email protected]>

Author: adobrzyn

tests/unit_tests/test_bucketing.py

@@ -63,7 +63,10 @@ def test_generate_decode_buckets():
     bs_bucket_config = [1, 32, 128]
     blocks_bucket_config = [128, 128, 2048]
     max_blocks = 1024
+    max_model_len = 131072
+    block_size = 128
     buckets = linear.generate_decode_buckets(bs_bucket_config,
-                                             blocks_bucket_config, max_blocks)
+                                             blocks_bucket_config, max_blocks,
+                                             max_model_len, block_size)
     assert len(buckets) == 72
     assert all(blocks <= max_blocks for _, _, blocks in buckets)

vllm_gaudi/extension/bucketing/linear.py

@@ -1,6 +1,7 @@
 import itertools
 import operator
 import os
+import math
 from dataclasses import dataclass, field
 from typing import List, Tuple