From d84b98b40f045d103c0106e9aa8aec17742a98fc Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Fri, 9 Aug 2024 12:47:39 +0200 Subject: [PATCH] Early exit on server too. --- .gitignore | 1 + server/text_generation_server/models/globals.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 0de8b848..bd9d9125 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,4 @@ server/exllama_kernels/exllama_kernels/exllama_ext_hip.cpp data/ load_tests/*.json +server/fbgemmm diff --git a/server/text_generation_server/models/globals.py b/server/text_generation_server/models/globals.py index 7e22a56c..ecd1ab32 100644 --- a/server/text_generation_server/models/globals.py +++ b/server/text_generation_server/models/globals.py @@ -6,6 +6,8 @@ from typing import Dict, Optional from text_generation_server.utils.log import log_master ATTENTION = os.getenv("ATTENTION", "paged") +_expected = {"paged", "flashdecoding", "flashinfer"} +assert ATTENTION in _expected, f"Attention is not valid {ATTENTION}, expected {_expected}" log_master(logger.info, f"Using Attention = {ATTENTION}") MEM_POOL = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None