mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
258c7540c0
Co-authored-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> open source f8c0381a2bc50ee2739c3d8c2be481b31e5f00bd (#2736) Co-authored-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Add note for blackwell (#2742) Update the docs to workaround the extra-index-url issue (#2744) update README.md (#2751) Fix github io pages (#2761) Update
20 lines
421 B
Python
20 lines
421 B
Python
from typing import List
|
|
|
|
|
|
def next_positive_power_of_2(x: int) -> int:
|
|
if x < 1:
|
|
return 1
|
|
|
|
return 1 << (x - 1).bit_length()
|
|
|
|
|
|
def get_power_of_2_num_tokens_buckets(max_num_tokens) -> List[int]:
|
|
max_num_tokens = next_positive_power_of_2(max_num_tokens)
|
|
num_token_buckets = []
|
|
m = 1
|
|
while m <= max_num_tokens:
|
|
num_token_buckets.append(m)
|
|
m *= 2
|
|
|
|
return num_token_buckets
|