Merge branch 'extract_outliers' into debug

1 files changed, 26 insertions, 0 deletions

diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index e7261bc..08c108c 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -1435,3 +1435,29 @@ def dequant_min_max(xq, A, B, SA, SB, dtype=torch.half):
     x *= SA[1]/127
     x +=offset
     return x.to(dtype)
+
+def extract_outliers(A, SA, idx):
+    shapeA = SA[0]
+    formatA = SA[1]
+    assert formatA in ['col_turing', 'col_ampere']
+    assert A.device.type == 'cuda'
+
+    out = torch.zeros((shapeA[0], idx.numel()), dtype=torch.int8, device=A.device)
+
+    idx_size = ct.c_int32(idx.numel())
+    rows = ct.c_int32(shapeA[0])
+    cols = ct.c_int32(shapeA[1])
+    ptrA = get_ptr(A)
+    ptrIdx = get_ptr(idx)
+    ptrOut = get_ptr(out)
+
+    if formatA == 'col_turing':
+        lib.cextractOutliers_turing(ptrA, ptrIdx, ptrOut, idx_size, rows, cols)
+    elif formatA == 'col_ampere':
+        lib.cextractOutliers_ampere(ptrA, ptrIdx, ptrOut, idx_size, rows, cols)
+
+    return out
+
+
+
+