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

2024-10-25 09:17:57 +02:00 · 2024-10-25 09:17:57 +02:00 · cece8635f8
parent 43df056eee
commit cece8635f8
2 changed files with 2 additions and 2 deletions
--- a/server/text_generation_server/layers/gptq/init.py
+++ b/server/text_generation_server/layers/gptq/init.py
@ -10,8 +10,8 @@ from text_generation_server.utils.weights import Weight, Weights, WeightsLoader

 if SYSTEM == "ipex":
    from .ipex import QuantLinear
-elif SYSTEM == "cuda":
-    from .cuda import QuantLinear
+elif SYSTEM in {"cuda", "rocm"}:
+    from .triton import QuantLinear


@dataclass
--- a/server/text_generation_server/layers/gptq/triton.py
+++ b/server/text_generation_server/layers/gptq/triton.py