From 8a7749b8fb678934a53ee72102644572a8d7b0a6 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Mon, 12 Aug 2024 17:59:37 -0400
Subject: [PATCH] fix: include create_exllama_buffers and set_device for
 exllama (#2407)

---
 server/text_generation_server/layers/gptq/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/server/text_generation_server/layers/gptq/__init__.py b/server/text_generation_server/layers/gptq/__init__.py
index 9c9b69d1..505caa59 100644
--- a/server/text_generation_server/layers/gptq/__init__.py
+++ b/server/text_generation_server/layers/gptq/__init__.py
@@ -422,12 +422,16 @@ elif CAN_EXLLAMA:
         if V2:
             from text_generation_server.layers.gptq.exllamav2 import (
                 QuantLinear as ExllamaQuantLinear,  # noqa: F401
+                create_exllama_buffers,  # noqa: F401
+                set_device,  # noqa: F401
             )
 
             HAS_EXLLAMA = "2"
         else:
             from text_generation_server.layers.gptq.exllama import (
                 Ex4bitLinear as ExllamaQuantLinear,  # noqa: F401
+                create_exllama_buffers,  # noqa: F401
+                set_device,  # noqa: F401
             )
 
             HAS_EXLLAMA = "1"