huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Black

Browse Source
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
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
server/text_generation_server/layers/attention/kv_cache.py
View File

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