Fixing rocm gptq by using triton code too (renamed cuda into triton). (…

…#2691)
huggingface · Oct 25, 2024 · cece863 · cece863
1 parent 43df056
commit cece863
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/server/text_generation_server/layers/gptq/__init__.py b/server/text_generation_server/layers/gptq/__init__.py
@@ -10,8 +10,8 @@
 
 if SYSTEM == "ipex":
     from .ipex import QuantLinear
-elif SYSTEM == "cuda":
-    from .cuda import QuantLinear
+elif SYSTEM in {"cuda", "rocm"}:
+    from .triton import QuantLinear
 
 
 @dataclass

diff --git a/...ext_generation_server/layers/gptq/cuda.py → ...t_generation_server/layers/gptq/triton.py b/...ext_generation_server/layers/gptq/cuda.py → ...t_generation_server/layers/gptq/triton.py