mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-22 15:32:08 +00:00
2358c2bb54
* Add basic FP8 KV cache support This change adds rudimentary FP8 KV cache support. The support is enabled by passing `--kv-cache-dtype fp8_e5m2` to the launcher. Doing so uses this type for the KV cache. However support is still limited: * Only the `fp8_e5m2` type is supported. * The KV cache layout is the same as `float16`/`bfloat16` (HND). * The FP8 KV cache is only supported for FlashInfer. * Loading of scales is not yet supported. * Fix Cargo.toml
48 lines
1.1 KiB
Python
48 lines
1.1 KiB
Python
import os
|
|
|
|
from text_generation_server.utils.import_utils import SYSTEM
|
|
|
|
from .common import Seqlen
|
|
|
|
if os.getenv("USE_FLASH_ATTENTION", "").lower() == "false":
|
|
raise ImportError("`USE_FLASH_ATTENTION` is false.")
|
|
if SYSTEM == "cuda":
|
|
from .cuda import (
|
|
PREFILL_IN_KV_CACHE,
|
|
SUPPORTS_WINDOWING,
|
|
attention,
|
|
paged_attention,
|
|
reshape_and_cache,
|
|
)
|
|
elif SYSTEM == "rocm":
|
|
from .rocm import (
|
|
PREFILL_IN_KV_CACHE,
|
|
SUPPORTS_WINDOWING,
|
|
attention,
|
|
paged_attention,
|
|
reshape_and_cache,
|
|
)
|
|
elif SYSTEM == "ipex":
|
|
from .ipex import (
|
|
PREFILL_IN_KV_CACHE,
|
|
SUPPORTS_WINDOWING,
|
|
attention,
|
|
paged_attention,
|
|
reshape_and_cache,
|
|
)
|
|
else:
|
|
raise ImportError(f"System {SYSTEM} doesn't support flash/paged attention")
|
|
|
|
# KVCache needs `reshape_and_cache`, so ensure that it is defined already.
|
|
from .kv_cache import KVCache
|
|
|
|
__all__ = [
|
|
"attention",
|
|
"paged_attention",
|
|
"reshape_and_cache",
|
|
"PREFILL_IN_KV_CACHE",
|
|
"SUPPORTS_WINDOWING",
|
|
"KVCache",
|
|
"Seqlen",
|
|
]
|