under test, not sure no errors

ixformer_sdk/inference/functions/bnb_qgemm.py

@@ -0,0 +1,56 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+
+__all__ = ["bnb_qgemm", "ref_bnb_qgemm"]
+
+
+# qA : quant input              shape : [bs, in_feature]
+# qW : quant weight             shape : [out_feature, in_feature]
+# SA : scale vector of qA       shape : [bs]
+# SW : scale vector of qW       shape : [out_feature]
+
+
+def ref_bnb_qgemm(
+    qA: torch.Tensor,
+    qW: torch.Tensor,
+    SA: torch.Tensor,
+    SW: torch.Tensor,
+    training: bool = False,
+    scaleA: float = 127.0,
+    scaleW: float = 127.0,
+):
+    y = torch.nn.functional.linear(qA.to(torch.float), qW.to(torch.float))
+    out = torch.empty(y.shape, dtype = SA.dtype, device = SA.device)
+    for i in range(qA.size(0)):
+        for j in range(qW.size(0)): 
+            out[i][j] = y[i][j] * (SA[i].to(torch.float) / scaleA) * (SW[j].to(torch.float) / scaleW)
+    return out.to(SA.dtype)
+
+
+def bnb_qgemm(
+    qA: torch.Tensor,
+    qW: torch.Tensor,
+    SA: torch.Tensor,
+    SW: torch.Tensor,
+    training: bool = False,
+    scaleA: float = 127.0,
+    scaleW: float = 127.0,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        qA:                 (bs, in_feature)                torch.int8
+        qW:                 (out_feature, in_feature)       torch.int8
+        SA:                 (bs)                            torch.half
+            scale vector of qA 
+        SA:                 (out_feature)                   torch.half
+            scale vector of qW 
+        training:                                           bool
+        scaleA:                                             float
+        scaleW:                                             float
+    Returns:
+        Tensor:             (bs, out_feature)               torch.half
+    """
+    return ops.infer.bnb_qgemm(qA, qW, SA, SW, scaleA, scaleW)