From f3bc03843003f7ad8f1aa8e85bd29b85264e2b2a Mon Sep 17 00:00:00 2001 From: Mohit Sharma Date: Wed, 11 Sep 2024 06:52:30 +0000 Subject: [PATCH] style --- server/text_generation_server/models/flash_causal_lm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py index 7a3f57ab..bb856078 100644 --- a/server/text_generation_server/models/flash_causal_lm.py +++ b/server/text_generation_server/models/flash_causal_lm.py @@ -1162,6 +1162,10 @@ class FlashCausalLM(Model): f"PyTorch TunableOp is enabled. The warmup may take several minutes, picking the ROCm optimal matrix multiplication kernel for the target lengths {', '.join([str(seqlen) for seqlen in tuning_sequences])}, with typical 5-8% latency improvement for small sequence lengths. The picked GEMMs are saved in the file {tunableop_filepath}. To disable TunableOp, please launch TGI with `PYTORCH_TUNABLEOP_ENABLED=0`.", ) + torch.cuda.tunable.set_filename( + tunableop_filepath, insert_device_ordinal=False + ) + if os.path.isfile(tunableop_filepath): log_master( logger.info,