Fixup some imports

huggingface · Feb 4, 2025 · 27decc5 · 27decc5
1 parent 17b5f57
commit 27decc5
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 12 deletions.
diff --git a/server/text_generation_server/layers/compressed_tensors/w8a8_int.py b/server/text_generation_server/layers/compressed_tensors/w8a8_int.py
@@ -6,15 +6,16 @@
 from compressed_tensors.quantization import QuantizationArgs, QuantizationType
 
 from text_generation_server.layers.fp8 import _load_scalar_or_matrix_scale
+from text_generation_server.utils.import_utils import SYSTEM
 from text_generation_server.utils.kernels import load_kernel
 from text_generation_server.utils.log import log_once
 from text_generation_server.utils.weights import Weight, Weights, WeightsLoader
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
 

diff --git a/server/text_generation_server/layers/fp8.py b/server/text_generation_server/layers/fp8.py
@@ -15,14 +15,15 @@
 )
 from text_generation_server.utils.log import log_once
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
 try:
+    # TODO: needs to be ported over to MoE and used on CUDA.
     from moe_kernels.fp8_utils import w8a8_block_fp8_matmul, per_token_group_quant_fp8
 except ImportError:
     w8a8_block_fp8_matmul = None

diff --git a/server/text_generation_server/layers/marlin/fp8.py b/server/text_generation_server/layers/marlin/fp8.py
@@ -8,13 +8,14 @@
     _check_marlin_kernels,
     permute_scales,
 )
+from text_generation_server.utils.import_utils import SYSTEM
 from text_generation_server.utils.kernels import load_kernel
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
 

diff --git a/server/text_generation_server/layers/marlin/gptq.py b/server/text_generation_server/layers/marlin/gptq.py
@@ -16,13 +16,14 @@
 from text_generation_server.utils.log import log_once
 from text_generation_server.utils.weights import Weight, Weights, WeightsLoader
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
+
 try:
     major, _minor = torch.cuda.get_device_capability()
     has_sm_8_0 = major >= 8

diff --git a/server/text_generation_server/layers/marlin/marlin.py b/server/text_generation_server/layers/marlin/marlin.py
@@ -5,14 +5,15 @@
 import torch.nn as nn
 
 from text_generation_server.layers.marlin.util import _check_marlin_kernels
+from text_generation_server.utils.import_utils import SYSTEM
 from text_generation_server.utils.kernels import load_kernel
 from text_generation_server.utils.weights import Weight, Weights, WeightsLoader
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
 

diff --git a/server/text_generation_server/layers/marlin/util.py b/server/text_generation_server/layers/marlin/util.py
@@ -6,11 +6,11 @@
 from text_generation_server.utils.import_utils import SYSTEM
 from text_generation_server.utils.kernels import load_kernel
 
-try:
+if SYSTEM == "cuda":
     marlin_kernels = load_kernel(
         module="quantization", repo_id="kernels-community/quantization"
     )
-except ImportError:
+else:
     marlin_kernels = None
 
 try: