修改awq workspace 申请

vllm/_custom_ops.py

@@ -143,6 +143,12 @@ def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
 
 # quantization ops
 # awq
+def GetAWQShareWorkspaceSize()->int:
+    return quant_ops.GetAWQShareWorkspaceSize()
+
+def GetAWQShareWorkspace()->torch.Tensor:
+    return quant_ops.GetAWQShareWorkspace()
+
 def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor,
                    zeros: torch.Tensor, split_k_iters: int, thx: int,
                    thy: int) -> torch.Tensor:

vllm/model_executor/layers/quantization/awq.py

@@ -20,12 +20,8 @@ class AWQShareWorkSpace:
         return cls._instance
 
     def _initialize(self):
-        self.awqworkshapcesize = 2 << 29
-        self.awqworkshapce = torch.zeros(self.awqworkshapcesize // 2 + 1, dtype=torch.float16).cuda()
-        #print("AWQShareWorkSpace _initialize\n")
-        #print("self.awqworkshapce.device:",self.awqworkshapce.device)
-
-
+        self.awqworkshapcesize = ops.GetAWQShareWorkspaceSize()
+        self.awqworkshapce = ops.GetAWQShareWorkspace()
 
 class AWQConfig(QuantizationConfig):
     """Config class for AWQ.
@@ -202,6 +198,7 @@ class AWQLinearMethod(LinearMethodBase):
         else:
             padding_group=0
             
+        if m<4096: 
             out = ops.awq_gemm(reshaped_x,
                             qweight,
                             zeros_and_scales,
@@ -212,14 +209,15 @@ class AWQLinearMethod(LinearMethodBase):
                             padding_group,
                             self.awqsingleton.awqworkshapce,
                             self.awqsingleton.awqworkshapcesize)
+        else: 
             #下面是采用rocblas的做法
-        # deqweight=ops.dequant_w4_gemm_colmajor(    #shape[n,k/8]--->[n,k]
-        #                   qweight, 
-        #                   zeros_and_scales,
-        #                   k,
-        #                   n,
-        #                   self.quant_config.group_size)
-        # output=F.linear(reshaped_x, deqweight)    
+            deqweight=ops.dequant_w4_gemm_colmajor(    #shape[n,k/8]--->[n,k]
+                            qweight, 
+                            zeros_and_scales,
+                            k,
+                            n,
+                            self.quant_config.group_size)
+            output=F.linear(reshaped_x, deqweight[:,0:k])    
         
         if bias is not None:
             out.add_(bias)