biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[TPU][V1] Fix exponential padding when max-num-batched-tokens is not a power of 2 (#16596)

Browse Source 
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
Nicolò Lucchesi
2025-04-14 19:01:05 +02:00
committed by GitHub
parent aa29841ede
commit b3f2fddd17
2 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/v1/tpu/worker/test_tpu_model_runner.py
View File

@@ -299,6 +299,18 @@ def test_get_paddings():
actual_paddings = _get_token_paddings(min_token_size, max_token_size,
padding_gap)
assert actual_paddings == expected_paddings
# Exponential padding.
max_token_size, padding_gap = 1024, 0
expected_paddings = [16, 32, 64, 128, 256, 512, 1024]
actual_paddings = _get_token_paddings(min_token_size, max_token_size,
padding_gap)
assert actual_paddings == expected_paddings
# Exponential padding with max_token_size not a power of two.
max_token_size = 317
expected_paddings = [16, 32, 64, 128, 256, 512]
actual_paddings = _get_token_paddings(min_token_size, max_token_size,
padding_gap)
assert actual_paddings == expected_paddings
def test_get_padded_token_len():