From 6d6b0bdcc4d8d65e53b1f28c0c2a41b9fe75faf1 Mon Sep 17 00:00:00 2001
From: Mohit Sharma <mohit21sharma.ms@gmail.com>
Date: Tue, 2 Jul 2024 13:08:56 +0000
Subject: [PATCH] fix formatting

---
 server/text_generation_server/cli.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py
index 9ff72569..8dd1e6e8 100644
--- a/server/text_generation_server/cli.py
+++ b/server/text_generation_server/cli.py
@@ -29,6 +29,12 @@ class Dtype(str, Enum):
     bloat16 = "bfloat16"
 
 
+class KVDtype(str, Enum):
+    auto = "auto"
+    fp8 = "fp8"
+    fp8_e5m2 = "fp8_e5m2"
+
+
 @app.command()
 def serve(
     model_id: str,
@@ -37,7 +43,7 @@ def serve(
     quantize: Optional[Quantization] = None,
     speculate: Optional[int] = None,
     dtype: Optional[Dtype] = None,
-    kv_cache_dtype: str = "auto",
+    kv_cache_dtype: KVDtype = "auto",
     trust_remote_code: bool = False,
     uds_path: Path = "/tmp/text-generation-server",
     logger_level: str = "INFO",