Apply suggestion on GPTQ buffer setup

Signed-off-by: cyang49 <[email protected]>
Signed-off-by: Chih-Chieh-Yang <[email protected]>

Author: cyang49

server/text_generation_server/server.py

@@ -278,31 +278,29 @@ async def serve_inner(
 
         if quantize == "gptq" and deployment_framework == "tgis_native":
             from text_generation_server.utils.layers import HAS_GPTQ_CUDA, EXLLAMA_VERSION
-            if HAS_GPTQ_CUDA:
-                if EXLLAMA_VERSION is not None:
-                    try:
-                        # When using GPTQ, Exllama kernels need some global kernels
-                        # For which we have the final shapes only after the model has loaded
-                        # This will allocate those buffers.
-
-                        if EXLLAMA_VERSION == "1":
-                            from text_generation_server.utils.gptq.exllama import (
-                                create_exllama_buffers, set_device,
-                            )
-                            set_device(device)
-                            create_exllama_buffers(max_sequence_length)
-                        else:
-                            assert EXLLAMA_VERSION == "2"
-                            from text_generation_server.utils.gptq.exllamav2 import (
-                                set_device, Ex4bitLinearV2,
-                            )
-                            set_device(device)
-                            for _, submodule in model.model.named_modules():
-                                if isinstance(submodule, Ex4bitLinearV2):
-                                    submodule.post_init()  # make q matrix and set scratch space
-
-                    except ImportError:
-                        print("WARN: Error setting up GPTQ exllama buffers")
+            if HAS_GPTQ_CUDA and EXLLAMA_VERSION is not None:
+                try:
+                    # When using GPTQ, Exllama kernels need some global kernels
+                    # For which we have the final shapes only after the model has loaded
+                    # This will allocate those buffers.
+                    if EXLLAMA_VERSION == "1":
+                        from text_generation_server.utils.gptq.exllama import (
+                            create_exllama_buffers, set_device,
+                        )
+                        set_device(device)
+                        create_exllama_buffers(max_sequence_length)
+                    elif EXLLAMA_VERSION == "2":
+                        from text_generation_server.utils.gptq.exllamav2 import (
+                            set_device, Ex4bitLinearV2,
+                        )
+                        set_device(device)
+                        for _, submodule in model.model.named_modules():
+                            if isinstance(submodule, Ex4bitLinearV2):
+                                submodule.post_init()  # make q matrix and set scratch space
+                    else:
+                        raise ValueError(f"Unsupported {EXLLAMA_VERSION=}") 
+                except ImportError:
+                    print("WARN: Error setting up GPTQ exllama buffers")
 
         if local_rank == 0 and device.type == "cuda":
             # Log GPU memory stats at startup