A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[TPU][V1] Fix exponential padding when `max-num-batched-tokens` is not a power of 2 (#16596)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 14, 2025 at 19:01 UTC
b3f2fddd172cef23b42ffaf6c226877b6588964c
2 files changed
+15
-1
tests/v1/tpu/worker/test_tpu_model_runner.py
+12
@@ -299,6 +299,18 @@ def test_get_paddings():
299
actual_paddings = _get_token_paddings(min_token_size, max_token_size,
300
padding_gap)
301
assert actual_paddings == expected_paddings
302
+ # Exponential padding.
303
+ max_token_size, padding_gap = 1024, 0
304
+ expected_paddings = [16, 32, 64, 128, 256, 512, 1024]
305
+ actual_paddings = _get_token_paddings(min_token_size, max_token_size,
306
+ padding_gap)
307
+ assert actual_paddings == expected_paddings
308
+ # Exponential padding with max_token_size not a power of two.
309
+ max_token_size = 317
310
+ expected_paddings = [16, 32, 64, 128, 256, 512]
311
+ actual_paddings = _get_token_paddings(min_token_size, max_token_size,
312
+ padding_gap)
313
+ assert actual_paddings == expected_paddings
314
315
316
def test_get_padded_token_len():
vllm/v1/worker/tpu_model_runner.py
+3
-1
@@ -1040,9 +1040,11 @@ def _get_token_paddings(min_token_size: int, max_token_size: int,
1040
1041
if padding_gap == 0:
1042
logger.info("Using exponential token paddings:")
1043
- while num <= max_token_size:
1043
+ while True:
1044
logger.info(" %d", num)
1045
paddings.append(num)
1046
+ if num >= max_token_size:
1047
+ break
1048
num *= 2
1049
1050
else: