Early exit on server too.

huggingface · Aug 9, 2024 · af49d54 · af49d54
1 parent 3e607a7
commit af49d54
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git 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
@@ -6,6 +6,8 @@
 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
Original file line number	Diff line number	Diff line change
Expand Up		@@ -18,3 +18,4 @@ server/exllama_kernels/exllama_kernels/exllama_ext_hip.cpp

		data/
		load_tests/*.json
		server/fbgemmm