huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-24 00:12:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix bug

Browse Source
This commit is contained in:
fxmarty 2024-05-17 16:03:15 +00:00
parent 422bf1f986
commit cd3c28cfe7
2 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
server/text_generation_server/models/flash_causal_lm.py
View File

@ -827,7 +827,7 @@ class FlashCausalLM(Model):
self.device,
)
if SYSTEM == "rocm":
if SYSTEM == "rocm" and self.speculate is None or self.speculate == 0:
if (
os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
@ -875,7 +875,11 @@ class FlashCausalLM(Model):
logger.info(f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}")
# Warmup cuda graphs
for bs in CUDA_GRAPHS:
if self.speculate is None or self.speculate + 1 <= bs:
if (
self.speculate is None
or self.speculate == 0
or self.speculate + 1 <= bs
):
self.cuda_graph_warmup(bs, max_s, max_bt)
except torch.cuda.OutOfMemoryError:
logger.exception(f"Decode cuda graph warmup failed")