Added warning message for failed awq_ext import to enhance error clarity

awq/modules/fused/mlp.py

@@ -2,13 +2,14 @@
 import torch.nn.functional as F
 from awq.modules.linear.gemm import WQLinear_GEMM
 from awq.modules.linear.gemv import WQLinear_GEMV
+import warnings
 
 try:
     import awq_ext  # with CUDA kernels
-
     AWQ_INSTALLED = True
-except:
+except Exception as e:
     AWQ_INSTALLED = False
+    warnings.warn(f"AWQ extension could not be imported. Error: {e}")
 
 
 class QuantFusedMLP(nn.Module):

awq/modules/fused/moe.py

@@ -1,12 +1,13 @@
 import torch
 from typing import Dict
+import warnings
 
 try:
     import awq_ext  # with CUDA kernels
-
     AWQ_INSTALLED = True
-except:
+except Exception as e:
     AWQ_INSTALLED = False
+    warnings.warn(f"AWQ extension could not be imported. Error: {e}")
 
 
 class FusedSparseMoeBlock(torch.nn.Module):

awq/modules/fused/norm.py

@@ -1,13 +1,13 @@
 import torch
 from torch import nn
+import warnings
 
 try:
     import awq_ext  # with CUDA kernels
-
     AWQ_INSTALLED = True
-except:
+except Exception as e:
     AWQ_INSTALLED = False
-
+    warnings.warn(f"AWQ extension could not be imported. Error: {e}")
 
 class FasterTransformerRMSNorm(nn.Module):
     def __init__(self, weight, eps=1e-6):

awq/modules/linear/gemm.py

@@ -3,13 +3,14 @@
 from torch.autograd import Function
 from awq.utils.utils import get_best_device
 from awq.utils.packing_utils import dequantize_gemm
+import warnings
 
 try:
-    import awq_ext  # with CUDA kernels (AutoAWQ_kernels)
-
+    import awq_ext  # with CUDA kernels
     AWQ_INSTALLED = True
-except:
+except Exception as e:
     AWQ_INSTALLED = False
+    warnings.warn(f"AWQ extension could not be imported. Error: {e}")
 
 
 # Adapted from https://github.com/compressa-ai/AutoAWQ/tree/dev

awq/modules/linear/gemv.py

@@ -1,13 +1,13 @@
 import torch
 import torch.nn as nn
+import warnings
 
 try:
     import awq_ext  # with CUDA kernels
-
     AWQ_INSTALLED = True
-except:
+except Exception as e:
     AWQ_INSTALLED = False
-
+    warnings.warn(f"AWQ extension could not be imported. Error: {e}")
 
 def make_divisible(c, divisor):
     return (c + divisor - 1) // divisor