[Misc] Refactor get_kv_cache_spec into AttentionLayerBase (#26587)
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
@@ -137,6 +137,15 @@ def set_default_torch_num_threads(num_threads: int):
|
||||
torch.set_num_threads(old_num_threads)
|
||||
|
||||
|
||||
def kv_cache_dtype_str_to_dtype(
|
||||
kv_cache_dtype: str, model_config: ModelConfig
|
||||
) -> torch.dtype:
|
||||
if kv_cache_dtype == "auto":
|
||||
# Model config may not be specified for unit tests, default to float16
|
||||
return model_config.dtype if model_config else torch.half
|
||||
return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user