[Misc] Refactor get_kv_cache_spec into AttentionLayerBase (#26587)

Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
Nicolò Lucchesi
2025-10-18 15:51:21 +02:00
committed by GitHub
parent ab4be40fc5
commit b26b70bec4
10 changed files with 151 additions and 118 deletions

View File

@@ -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")