This commit is contained in:
Nicolas Patry 2024-10-08 09:28:07 +02:00
parent a75c9a21e8
commit 53f9b18086
No known key found for this signature in database
GPG Key ID: 64AF4752B2967863

View File

@ -24,10 +24,8 @@ class KVCache:
): ):
"""Construct the key-value cache for a layer.""" """Construct the key-value cache for a layer."""
if ( if dtype == torch.float8_e5m2 and (
dtype == torch.float8_e5m2 ATTENTION != "flashinfer" or SYSTEM != "cuda"
and (ATTENTION != "flashinfer"
or SYSTEM != "cuda")
): ):
raise ValueError( raise ValueError(
"float8_e5m2 KV cache is currently only supported for flashinfer on CUDA" "float8_e5m2 KV cache is currently only supported for flashinfer on CUDA"