File size: 1,041 Bytes
2976eec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | # SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
from typing import List
import torch
import triton
import json
def prev_power_of_2(x: int) -> int:
out = triton.next_power_of_2(x)
return out // 2 if out > x else out
STATIC_MAX_SEQ_LENS: List[int] = []
USE_RUNTIME_MAX_SEQ_LEN: bool = False
def autotune_max_seq_len(runtime_max_seq_len: int) -> int:
global USE_RUNTIME_MAX_SEQ_LEN
if USE_RUNTIME_MAX_SEQ_LEN:
return prev_power_of_2(runtime_max_seq_len)
else:
if STATIC_MAX_SEQ_LENS == []:
return 1
for max_len in STATIC_MAX_SEQ_LENS:
if max_len >= runtime_max_seq_len:
return max_len
return STATIC_MAX_SEQ_LENS[-1]
def switch_to_contiguous_if_needed(x: torch.Tensor) -> torch.Tensor:
if x.stride(-1) == 1:
return x
return x.contiguous()
def serialize_dict(d: dict) -> str:
return json.dumps(d)
def deserialize_str(s: str) -> dict:
return json.loads(s)
|