fix: include create_exllama_buffers and set_device for exllama (#2407)

8a7749b8 · drbh · GitHub · 9a7830bd · 8a7749b8
Unverified Commit 8a7749b8 authored Aug 12, 2024 by drbh Committed by GitHub Aug 12, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

server/text_generation_server/layers/gptq/__init__.py server/text_generation_server/layers/gptq/__init__.py +4 -0

No files found.
--- 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"