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

Make check more obvious

Browse Source
This commit is contained in:
Daniël de Kok 2024-10-16 13:54:57 +00:00
parent aa92e451a0
commit 751f1bb815
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.itemsize == 1
and dtype.is_floating_point
and (ATTENTION != "flashinfer" or SYSTEM != "cuda")
if dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and (
ATTENTION != "flashinfer" or SYSTEM != "cuda"
):
raise ValueError(
"FP8 KV cache is currently only supported for flashinfer on CUDA"