diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py index acd97f45..20f93820 100644 --- a/server/text_generation_server/models/flash_causal_lm.py +++ b/server/text_generation_server/models/flash_causal_lm.py @@ -802,7 +802,7 @@ class FlashCausalLM(Model): try: logger.info("Experimental support for Cuda Graphs is enabled") # Warmup cuda graphs - for bs in [1, 2, 4] + [8 * i for i in range(8)]: + for bs in [1, 2, 4] + [8 * i for i in range(1, 9)]: if self.speculate is None or self.speculate + 1 <= bs: self.cuda_graph_warmup(bs, max_s, max_bt) except Exception: