Most tests passing.

4 files changed, 2446 insertions, 152 deletions

diff --git a/tests/test_autograd.py b/tests/test_autograd.py
new file mode 100644
index 0000000..d2b5d59
--- /dev/null
+++ b/tests/test_autograd.py
@@ -0,0 +1,270 @@
+import pytest
+
+import torch
+import bitsandbytes as bnb
+
+from itertools import product
+
+n = 1
+k = 25
+dim1 = torch.randint(16,64, size=(n,)).tolist()
+dim2 = torch.randint(32,96, size=(n,)).tolist()
+dim3 = torch.randint(32,96, size=(n,)).tolist()
+dim4 = torch.randint(32,96, size=(n,)).tolist()
+funcs = [(torch.bmm, bnb.bmm_cublas), (torch.matmul, bnb.matmul_cublas)]
+str_funcs = ['bmm', 'matmul']
+req_grad = [(False, False), (True, False), (True, True), (False, True)]
+req_grad_str = ['FF', 'TF', 'TT', 'FT']
+transpose = [(False, False), (False, True), (True, True), (True, False)]
+str_transpose = ['FF', 'FT', 'TT', 'TF']
+dtype = [torch.float32, torch.float16]
+values = list(product(dim1,dim2,dim3,dim4,funcs, dtype, req_grad, transpose))
+str_values = list(product(dim1,dim2,dim3,dim4,str_funcs, dtype, req_grad_str, str_transpose))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}'.format(*vals) for vals in str_values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose", values, ids=names)
+def test_matmul(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose):
+    dim2 = dim2 - (dim2 % 16)
+    dim3 = dim3 - (dim3 % 16)
+    dim4 = dim4 - (dim4 % 16)
+    for i in range(k):
+
+        # normal multiply
+        if funcs[0] in [torch.mm, torch.matmul]:
+            dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2)
+            dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3)
+            A = torch.randn(size=dimA, device='cuda', requires_grad=req_grad[0])
+            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1])
+            target = torch.randn(size=(dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            torch.nn.init.xavier_uniform_(B)
+
+            if not transpose[0] and not transpose[1]:
+                out_torch = funcs[0](A, B)
+                out_bnb = funcs[1](A, B)
+            elif not transpose[0] and transpose[1]:
+                out_torch = funcs[0](A, B.t())
+                out_bnb = funcs[1](A, B.t())
+            elif transpose[0] and not transpose[1]:
+                out_torch = funcs[0](A.t(), B)
+                out_bnb = funcs[1](A.t(), B)
+            elif transpose[0] and transpose[1]:
+                out_torch = funcs[0](A.t(), B.t())
+                out_bnb = funcs[1](A.t(), B.t())
+
+            n = out_bnb.numel()
+            idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
+            assert (idx==0).sum().item() < n*0.0175
+            idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
+            assert (idx==0).sum().item() < n*0.001
+
+            if any(req_grad):
+                out_bnb.data.copy_(out_torch)
+                torch.cuda.synchronize()
+                loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean()
+                loss_bnb.backward()
+                gradA1 = A.grad
+                gradB1 = B.grad
+                A.grad = None
+                B.grad = None
+
+                loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean()
+                loss_torch.backward()
+                gradA2 = A.grad
+                gradB2 = B.grad
+                A.grad = None
+                B.grad = None
+
+            if req_grad[0]:
+                torch.testing.assert_allclose(gradA1, gradA2, atol=0.015, rtol=0.1)
+            if req_grad[1]:
+                n = gradB1.numel()
+                idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.1
+                idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.02
+                torch.testing.assert_allclose(gradB1, gradB2, atol=0.18, rtol=0.3)
+
+        # batched matrix multiply
+        if funcs[0] in [torch.bmm, torch.matmul]:
+            A = torch.randn(size=(dim1, dim2, dim3), device='cuda', requires_grad=req_grad[0])
+            B = torch.randn(size=(dim1, dim3, dim4), device='cuda', requires_grad=req_grad[1])
+            target = torch.randn(size=(dim1, dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            torch.nn.init.xavier_uniform_(B)
+
+            out_torch = funcs[0](A, B)
+            out_bnb = funcs[1](A, B)
+
+            n = out_bnb.numel()
+            idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
+            assert (idx==0).sum().item() < n*0.01
+            torch.testing.assert_allclose(out_bnb, out_torch, atol=0.027, rtol=0.2)
+
+            if any(req_grad):
+                out_bnb.data.copy_(out_torch)
+                torch.cuda.synchronize()
+                loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean()
+                loss_bnb.backward()
+                gradA1 = A.grad
+                gradB1 = B.grad
+                A.grad = None
+                B.grad = None
+
+                loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean()
+                loss_torch.backward()
+                gradA2 = A.grad
+                gradB2 = B.grad
+                A.grad = None
+                B.grad = None
+
+            if req_grad[0]:
+                torch.testing.assert_allclose(gradA1, gradA2, atol=0.015, rtol=0.1)
+            if req_grad[1]:
+                n = gradB1.numel()
+                idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.1
+                idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.02
+
+        if funcs[0] in [torch.matmul]:
+            dim1 = dim1 - (dim1 % 16)
+            A = torch.randn(size=(dim1, dim2, dim3), device='cuda', requires_grad=req_grad[0])
+            dimB = (dim4, dim3) if transpose[1] else (dim3, dim4)
+            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1])
+            target = torch.randn(size=(dim1, dim2, dim4), device='cuda', requires_grad=req_grad[1])
+            torch.nn.init.xavier_uniform_(B)
+
+            if transpose[1]:
+                out_torch = funcs[0](A, B.t())
+                out_bnb = funcs[1](A, B.t())
+            else:
+                out_torch = funcs[0](A, B)
+                out_bnb = funcs[1](A, B)
+
+            n = out_bnb.numel()
+            idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
+            assert (idx==0).sum().item() < n*0.0175
+            idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
+            assert (idx==0).sum().item() < n*0.001
+
+            if any(req_grad):
+                out_bnb.data.copy_(out_torch)
+                torch.cuda.synchronize()
+                loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean()
+                loss_bnb.backward()
+                gradA1 = A.grad
+                gradB1 = B.grad
+                A.grad = None
+                B.grad = None
+
+                loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean()
+                loss_torch.backward()
+                gradA2 = A.grad
+                gradB2 = B.grad
+                A.grad = None
+                B.grad = None
+
+            if req_grad[0]:
+                torch.testing.assert_allclose(gradA1, gradA2, atol=0.015, rtol=0.1)
+            if req_grad[1]:
+                n = gradB1.numel()
+                idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.1
+                idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
+                assert (idx==0).sum().item() < n*0.02
+
+
+n = 1
+k = 3
+dim1 = torch.randint(16,64, size=(n,)).tolist()
+dim2 = torch.randint(32,96, size=(n,)).tolist()
+dim3 = torch.randint(32,96, size=(n,)).tolist()
+dim4 = torch.randint(32,96, size=(n,)).tolist()
+
+#dim1 = (17,)
+#dim2 = (7,)
+#dim3 = (37,)
+#dim4 = (23,)
+
+decomp = [0.0, 6.0]
+funcs = [(torch.matmul, bnb.matmul)]
+str_funcs = ['matmul']
+req_grad = [(False, False), (True, False), (True, True), (False, True)]
+req_grad_str = ['FF', 'TF', 'TT', 'FT']
+transpose = [(False, True), (False, False)]
+str_transpose = ['NT', 'NN']
+dtype = [torch.float16]
+has_fp16_weights = [True, False]
+values = list(product(dim1,dim2,dim3,dim4,funcs, dtype, req_grad, transpose, decomp, has_fp16_weights))
+str_values = list(product(dim1,dim2,dim3,dim4,str_funcs, dtype, req_grad_str, str_transpose, decomp, has_fp16_weights))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_func_{4}_dtype_{5}_requires_grad_{6}_transpose_{7}_decomp_{8}_has_fp16_weights_{9}'.format(*vals) for vals in str_values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights", values, ids=names)
+def test_matmullt(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights):
+    dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2)
+    dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3)
+    outlier_dim = torch.randint(0, dimA[1], size=(dimA[1]//8,), device='cuda')
+
+    for i in range(k):
+
+        # normal multiply
+        if funcs[0] in [torch.mm, torch.matmul]:
+            A = torch.randn(size=dimA, device='cuda', requires_grad=req_grad[0], dtype=dtype)
+            if decomp == 6.0:
+                with torch.no_grad():
+                    A[:, outlier_dim] = 6.0
+            B = torch.randn(size=dimB, device='cuda', requires_grad=req_grad[1], dtype=dtype)
+            target = torch.randn(size=(dim2, dim4), device='cuda', requires_grad=req_grad[1], dtype=dtype)
+            torch.nn.init.xavier_uniform_(B)
+            B2 = B.clone()
+
+            state = bnb.MatmulLtState()
+            state.threshold = decomp
+            state.has_fp16_weights = has_fp16_weights
+            if not has_fp16_weights:
+                if not transpose[0] and not transpose[1]: B2 = B2.t().contiguous()
+                state.CB, CBt, state.SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B2)
+                B2 = state.CB
+
+            if not transpose[0] and transpose[1]:
+                out_torch = funcs[0](A, B.t())
+                out_bnb = funcs[1](A, B2, state=state)
+            elif not transpose[0] and not transpose[1]:
+                out_torch = funcs[0](A, B)
+                out_bnb = funcs[1](A, B2.t(), state=state)
+
+            n = out_bnb.numel()
+            err = torch.abs(out_bnb-out_torch).mean().item()
+            #print(f'abs error {err:.4f}')
+            idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1)
+            assert (idx==0).sum().item() < n*0.0175
+            idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2)
+            assert (idx==0).sum().item() < n*0.001
+
+            if has_fp16_weights:
+                if any(req_grad):
+                    out_bnb.data.copy_(out_torch)
+                    torch.cuda.synchronize()
+                    loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean()
+                    loss_bnb.backward()
+                    gradA1 = A.grad
+                    gradB1 = B.grad
+                    A.grad = None
+                    B.grad = None
+
+                    loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean()
+                    loss_torch.backward()
+                    gradA2 = A.grad
+                    gradB2 = B.grad
+                    A.grad = None
+                    B.grad = None
+
+                if req_grad[0]:
+                    torch.testing.assert_allclose(gradA1, gradA2, atol=0.015, rtol=0.1)
+                if req_grad[1]:
+                    n = gradB1.numel()
+                    assert torch.abs(gradB1).sum() > 0.0
+                    assert torch.abs(gradB2).sum() > 0.0
+                    idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3)
+                    assert (idx==0).sum().item() < n*0.1
+                    idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3)
+                    assert (idx==0).sum().item() < n*0.02
+                    torch.testing.assert_allclose(gradB1, gradB2, atol=0.18, rtol=0.3)
+
diff --git a/tests/test_functional.py b/tests/test_functional.py
index 2a7d308..6cbe58f 100644
--- a/tests/test_functional.py
+++ b/tests/test_functional.py
@@ -1,15 +1,76 @@
-# Copyright (c) Facebook, Inc. and its affiliates. 
-#   
-# This source code is licensed under the MIT license found in the 
-# LICENSE file in the root directory of this source tree.
 import pytest
+import math
+import random
+import time
 import torch
 import bitsandbytes as bnb
+import einops
 
 from itertools import product
 
 from bitsandbytes import functional as F
 
+torch.set_printoptions(precision=4, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000)
+k = 20
+
+def assert_all_approx_close(a, b, rtol, atol, count):
+    idx = torch.isclose(a, b, rtol, atol)
+    sumval = (idx==0).sum().item()
+    if sumval > count:
+        print(f'Too many values not close: assert {sumval} < {count}')
+        torch.testing.assert_allclose(a, b, rtol, atol)
+
+class FFN(torch.nn.Module):
+    def __init__(self, input_features, hidden_size, bias=True):
+        super(FFN, self).__init__()
+        self.fc1 = torch.nn.Linear(input_features, hidden_size, bias=bias)
+        self.fc2 = torch.nn.Linear(hidden_size, input_features, bias=bias)
+
+        with torch.no_grad():
+            torch.nn.init.xavier_uniform_(self.fc1.weight)
+            torch.nn.init.xavier_uniform_(self.fc2.weight)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+class Timer(object):
+    def __init__(self):
+        self.starts = {}
+        self.ends = {}
+        self.agg = {}
+
+    def tick(self, name='default'):
+        if name not in self.starts:
+            self.starts[name] = torch.cuda.Event(enable_timing=True)
+            self.ends[name] = torch.cuda.Event(enable_timing=True)
+            self.starts[name].record()
+        else:
+            ms = self.tock(name, evict=True, print_ms=False)
+
+    def tock(self, name='default', evict=True, print_ms=True):
+        if name in self.ends:
+            self.ends[name].record()
+            torch.cuda.synchronize()
+            ms = self.starts[name].elapsed_time(self.ends[name])
+            if name not in self.agg: self.agg[name] = 0.0
+            self.agg[name] += ms
+            if evict:
+                self.starts.pop(name)
+                self.ends.pop(name)
+
+        if print_ms and name in self.agg:
+            print('{0} took: {1:.5f}s'.format(name, self.agg[name]/1000.0))
+
+        return self.agg[name]
+
+    def reset(self):
+        self.starts  = {}
+        self.ends = {}
+        self.agg = {}
+        print('Resetting benchmark data')
+
 def setup():
     pass
 
@@ -64,8 +125,8 @@ def test_dynamic_quantization():
         diffs.append(diff.mean().item())
         reldiffs.append(reldiff.mean().item())
         assert diff.mean().item() < 0.0135
-    print(sum(diffs)/len(diffs))
-    print(sum(reldiffs)/len(reldiffs))
+    #print(sum(diffs)/len(diffs))
+    #print(sum(reldiffs)/len(reldiffs))
 
     for i in range(100):
         A1 = torch.rand(1024, 1024, device='cuda')
@@ -88,8 +149,8 @@ def test_dynamic_blockwise_quantization():
         diffs.append(diff.mean().item())
         reldiffs.append(reldiff.mean().item())
         assert diffs[-1] < 0.011
-    print(sum(diffs)/len(diffs))
-    print(sum(reldiffs)/len(reldiffs))
+    #print(sum(diffs)/len(diffs))
+    #print(sum(reldiffs)/len(reldiffs))
 
     diffs = []
     for i in range(100):
@@ -125,7 +186,7 @@ def test_percentile_clipping(gtype):
     n = 4
     step = 0
     percentile=5
-    for i in range(1000):
+    for i in range(k):
         step += 1
         g = torch.randn(n, n, dtype=gtype, device='cuda')
         gnorm1, clip2, gnorm_scale = F.percentile_clipping(g, gnorm_vec2, step, percentile=percentile)
@@ -145,69 +206,1653 @@ def test_percentile_clipping(gtype):
         torch.testing.assert_allclose(gnorm1, gnorm2)
 
 
+def quant(x):
+    max1 = torch.abs(x).max()
+    x = torch.round(x/max1*127)
+    return max1, x.to(torch.int8)
+
+def dequant(c, maxC):
+    return c.float()*(maxC/127)
+
+def mm_dequant(maxA, maxB, C):
+    return C.float()*(maxA/127)*(maxB/127)
+
+def quant_multi(x, dim):
+    max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True)
+    max1[max1==0] = 1.0
+    x = torch.round(x/max1*127)
+    return max1, x.to(torch.int8)
+
+def quant_multi_chunk(x, dim, chunk_size=32):
+    if dim==1:
+        x_chunked = einops.rearrange(x, '(c a) b -> c a b', c=chunk_size)
+        max1 = torch.amax(torch.abs(x_chunked), dim=dim+1, keepdim=True)
+        max1 = torch.tile(max1, (1, 1, x.shape[1]))
+        max1 = max1.view(x.shape)
+    elif dim==0:
+        x_chunked = einops.rearrange(x, 'a (b c) -> a b c', c=chunk_size)
+        max1 = torch.amax(torch.abs(x_chunked), dim=dim, keepdim=True)
+        max1 = torch.tile(max1, (x.shape[0], 1, 1))
+        max1 = max1.view(x.shape)
+    max1[max1==0] = 1.0
+    x = torch.round(x/max1*127)
+    return max1, x.to(torch.int8)
+
+def quant_minmax(A):
+    minA = A.min()
+    maxA = A.max()
+
+def mean(xx):
+    return sum(xx)/float(len(xx))
+
+#dim1 = torch.randint(1,1024*4, size=(4,)).tolist()
+#dim2 = torch.randint(1,1024*4, size=(4,)).tolist()
+dim1 = [1024*2]
+dim2 = [1024*16]
+methods = [(lambda x, dim: quant(x), lambda x, dim: quant(x), dequant, dequant, mm_dequant)]
+methods.append((quant_multi, quant_multi, dequant, dequant, mm_dequant))
+#methods.append((lambda x: quant_multi_chunk(x, dim=-1), lambda x: quant_multi_chunk(x, dim=0), dequant, dequant, mm_dequant))
+method_names = ['linear', 'vectorwise']
+batched = [False, True]
+values = list(product(dim1,dim2, methods, batched))
+values_names = list(product(dim1,dim2, method_names, batched))
+names = ['dim1_{0}_dim2_{1}_quant_{2}_batched_{3}'.format(*vals) for vals in values_names]
+@pytest.mark.parametrize("dim1, dim2, quant_methods, batched", values, ids=names)
+def test_approx_igemm(dim1, dim2, quant_methods, batched):
+    dim1 = dim1 - (dim1 % 32)
+    dim2 = dim2 - (dim2 % 32)
+    errors = []
+    relerrors = []
+    print('')
+    for i in range(5):
+        if batched:
+            A = torch.normal(0, 0.5, size=(32, dim1, dim2//32), device='cuda')
+            B = torch.normal(0, 0.5, size=(32, dim2//32, dim1), device='cuda')
+            maxA, Ac = quant_methods[0](A, 2)
+            maxB, Bc = quant_methods[1](B, 1)
+        else:
+            A = torch.normal(0, 0.5, size=(dim1, dim2), device='cuda')
+            B = torch.normal(0, 0.5, size=(dim2, dim1), device='cuda')
+            maxA, Ac = quant_methods[0](A, 1)
+            maxB, Bc = quant_methods[1](B, 0)
+        torch.testing.assert_allclose(quant_methods[2](maxA, Ac), A, atol=0.025, rtol=0.05)
+        if batched:
+            out2 = torch.bmm(A, B)
+            C = torch.bmm(Ac.float(), Bc.float())
+        else:
+            out2 = torch.mm(A, B)
+            C = F.igemm(Ac, Bc)
+        out = quant_methods[4](maxA, maxB, C)
+        std = out2.std()
+        out/= std
+        out2/= std
+        err = torch.abs(out-out2)
+        relerr = err/torch.abs(out2)
+        errors.append(err.mean().item())
+        relerrors.append(relerr.mean().item())
+    print(mean(errors))
+    print(mean(relerrors))
+
+
+
+
+
+
 def test_stable_embedding():
     layer = bnb.nn.StableEmbedding(1024, 1024)
     layer.reset_parameters()
 
 
-def test_dynamic_blockwise_quantization_cpu():
-    #A1 = torch.randn(1024, 1024, device='cpu')
-    #code = F.create_dynamic_map()
-    #for i in range(1000):
-    #    C, S = F.quantize_blockwise(A1, code=code)
-    #    A2 = F.dequantize_blockwise(C, S)
 
-    for i in range(10):
-        # equivalence with GPU blockwise quantization
-        A1 = torch.randn(1024, 1024, device='cpu')
-        C1, S1 = F.quantize_blockwise(A1)
-        C2, S2 = F.quantize_blockwise(A1.cuda())
-        torch.testing.assert_allclose(S1[0], S2[0].cpu())
-        # there seems to be some issues with precision in CUDA vs CPU
-        # not all elements are usually close, with couple off elements in a million
-        idx = torch.isclose(C1, C2.cpu())
-        assert (idx==0).sum().item() < 15
+n = 2
+hidden_dim = torch.randint(32,256, size=(n,)).tolist()
+batch_dim = torch.randint(16,256, size=(n,)).tolist()
+seq_dim = torch.randint(16,256, size=(n,)).tolist()
+transpose = [(False, False), (False, True), (True, False), (True, True)]
+values = list(product(hidden_dim,batch_dim, transpose, seq_dim))
+names = ['hidden_dim_{0}_batch_dim_{1},transpose_{2}_seq_dim_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("hidden_dim, batch_dim, transpose, seq_dim", values, ids=names)
+def test_igemm(hidden_dim, batch_dim, transpose, seq_dim):
+    hidden_dim = hidden_dim - (hidden_dim % 32)
+    batch_dim = batch_dim - (batch_dim % 16)
+    seq_dim = seq_dim - (seq_dim % 16)
+    for i in range(k):
+        shapeA = (batch_dim, hidden_dim) if not transpose[0] else (hidden_dim, batch_dim)
+        shapeB = ((32*random.randint(1, 4), hidden_dim) if transpose[1] else (hidden_dim, 32*random.randint(1, 4)))
+        A = torch.randint(-128, 127, size=shapeA, device='cuda').to(torch.int8)
+        B = torch.randint(-128, 127, size=shapeB, device='cuda').to(torch.int8)
+        if not transpose[0] and not transpose[1]:
+            out2 = torch.matmul(A.float(), B.float())
+            out = F.igemm(A, B)
+        elif not transpose[0] and transpose[1]:
+            out2 = torch.matmul(A.float(), B.t().float())
+            out = F.igemm(A, B.t())
+        elif transpose[0] and not transpose[1]:
+            out2 = torch.matmul(A.t().float(), B.float())
+            out = F.igemm(A.t(), B)
+        elif transpose[0] and transpose[1]:
+            out2 = torch.matmul(A.t().float(), B.t().float())
+            out = F.igemm(A.t(), B.t())
 
+        torch.testing.assert_allclose(out.float(), out2)
 
-    diffs = []
-    reldiffs = []
+    for i in range(k):
+        shapeA = (batch_dim, seq_dim, hidden_dim)
+        shapeB = ((32*random.randint(1, 4), hidden_dim) if transpose[1] else (hidden_dim, 32*random.randint(1, 4)))
+        A = torch.randint(-128, 127, size=shapeA, device='cuda').to(torch.int8)
+        B = torch.randint(-128, 127, size=shapeB, device='cuda').to(torch.int8)
+        if not transpose[0] and not transpose[1]:
+            out2 = torch.matmul(A.float(), B.float())
+            out = F.igemm(A, B)
+        elif not transpose[0] and transpose[1]:
+            out2 = torch.matmul(A.float(), B.t().float())
+            out = F.igemm(A, B.t())
+
+        torch.testing.assert_allclose(out.float(), out2)
+
+
+n = 3
+seq_dim = torch.randint(32,512, size=(n,)).tolist()
+hidden_dim = torch.randint(32,1024*4, size=(n,)).tolist()
+batch_dim = torch.randint(2,16, size=(n,)).tolist()
+values = list(product(seq_dim,hidden_dim,batch_dim))
+names = ['seq_dim{0}_hidden_dim{1}_batch_dim{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("seq_dim, hidden_dim, batch_dim", values, ids=names)
+def test_dim3_igemm(seq_dim, hidden_dim, batch_dim):
+    seq_dim = seq_dim - (seq_dim % 32)
+    hidden_dim = hidden_dim - (hidden_dim % 32)
+    batch_dim = batch_dim - (batch_dim % 2)
+    for i in range(25):
+        A = torch.randint(-128, 127, size=(batch_dim, seq_dim, hidden_dim), device='cuda').to(torch.int8)
+        B = torch.randint(-128, 127, size=(batch_dim, seq_dim, 1024), device='cuda').to(torch.int8)
+        out2 = torch.einsum('bsi, bso->io', A.float(), B.float())
+        iout = torch.empty(A.shape[2], B.shape[2], dtype=torch.int32, device=A.device)
+        out = F.igemm(A, B, out=iout)
+
+        torch.testing.assert_allclose(out.float(), out2)
+
+n = 2
+seq_dim = torch.randint(32,512, size=(n,)).tolist()
+hidden_dim = torch.randint(32,1024*4, size=(n,)).tolist()
+batch_dim = torch.randint(2,16, size=(n,)).tolist()
+transpose = [False, True]
+values = list(product(seq_dim,hidden_dim,batch_dim, transpose))
+names = ['seq_dim={0}_hidden_dim={1}_batch_dim={2}_transpose{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("seq_dim, hidden_dim, batch_dim, transpose", values, ids=names)
+def test_minmax_igemm(seq_dim, hidden_dim, batch_dim, transpose):
+
+    def min_max(x):
+        maxA = torch.amax(x, dim=2, keepdim=True)
+        minA = torch.amin(x, dim=2, keepdim=True)
+        scale = (maxA-minA)/2.0
+        return (127*(x-minA-scale)/scale).to(torch.int8), minA, scale
+
+    seq_dim = seq_dim - (seq_dim % 16)
+    hidden_dim = hidden_dim - (hidden_dim % 16)
+    batch_dim = batch_dim - (batch_dim % 2)
+    errs = []
+    relerrs = []
+    errs2 = []
+    relerrs2 = []
+    for i in range(k):
+        A = torch.normal(0.0, 0.5, size=(batch_dim, seq_dim, hidden_dim), device='cuda')
+        if transpose:
+            B = torch.normal(0, 0.5, size=(256, hidden_dim), device='cuda')
+        else:
+            B = torch.normal(0, 0.5, size=(hidden_dim, 256), device='cuda')
+        Ac, minA, scale = min_max(A)
+        if transpose:
+            maxB, Bc = quant_multi(B, dim=(1 if transpose else 0))
+            out = F.igemm(Ac, Bc.t())
+            out2 = torch.matmul(A,B.t())
+            offset = B.t().sum(0)*(minA+scale)
+            out = out.float()
+            out = (out*maxB.t()*scale/(127*127))+offset
+
+            maxA, Ac = quant_multi(A, dim=2)
+            out3 = F.igemm(Ac, Bc.t())
+            out3 = mm_dequant(maxA, maxB.t(), out3)
+        else:
+            maxB, Bc = quant_multi(B, dim=0)
+            offset = B.sum(0)*(minA+scale)
+            out = F.igemm(Ac, Bc)
+            out2 = torch.matmul(A,B)
+            out = out.float()
+            out = (out*maxB*scale/(127*127))+offset
+
+            maxA, Ac = quant_multi(A, dim=2)
+            out3 = F.igemm(Ac, Bc)
+            out3 = mm_dequant(maxA, maxB, out3)
+
+        std = out2.std()
+        out2 /= std
+        out /= std
+        out3 /= std
+
+        err = torch.abs(out-out2)
+        relerr = err/(torch.abs(out2)+1e-7)
+
+        err2 = torch.abs(out3-out2)
+        relerr2 = err2/(torch.abs(out2)+1e-7)
+
+        errs.append(err.mean().item())
+        relerrs.append(relerr.mean().item())
+        errs2.append(err2.mean().item())
+        relerrs2.append(relerr2.mean().item())
+    #print(mean(errs))
+    #print(mean(relerrs))
+    #print(mean(errs2))
+    #print(mean(relerrs2))
+    assert mean(errs) < 0.015
+    assert mean(relerrs) < 0.3
+
+n = 2
+dim1 = torch.randint(1,64, size=(n,)).tolist()
+dim2 = torch.randint(32,128, size=(n,)).tolist()
+dim3 = torch.randint(32,256, size=(n,)).tolist()
+dim4 = torch.randint(32,256, size=(n,)).tolist()
+transpose = [(False, False), (True, False), (False, True), (True, True)]
+values = list(product(dim1,dim2,dim3,dim4,transpose))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_transpose_{4}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dim4, transpose", values, ids=names)
+def test_ibmm(dim1, dim2, dim3, dim4, transpose):
+    dim2 = dim2 - (dim2 % 16)
+    dim3 = dim3 - (dim3 % 16)
+    dim4 = dim4 - (dim4 % 16)
+    for i in range(k):
+        shapeA = (dim1, dim3, dim2) if transpose[0] else (dim1, dim2, dim3)
+        shapeB = (dim1, dim4, dim3) if transpose[1] else (dim1, dim3, dim4)
+        A = torch.randint(-128, 127, size=shapeA, device='cuda').to(torch.int8)
+        B = torch.randint(-128, 127, size=shapeB, device='cuda').to(torch.int8)
+
+        if not transpose[0] and not transpose[1]:
+            out2 = torch.bmm(A.float(), B.float())
+            out = F.igemm(A, B)
+        elif not transpose[0] and transpose[1]:
+            out2 = torch.bmm(A.float(), B.permute([0, 2, 1]).float())
+            out = F.igemm(A, B.permute([0, 2, 1]))
+        elif transpose[0] and not transpose[1]:
+            out2 = torch.bmm(A.permute([0, 2, 1]).float(), B.float())
+            out = F.igemm(A.permute([0, 2, 1]), B)
+        elif transpose[0] and transpose[1]:
+            out2 = torch.bmm(A.permute([0, 2, 1]).float(), B.permute([0, 2, 1]).float())
+            out = F.igemm(A.permute([0, 2, 1]), B.permute([0, 2, 1]))
+        torch.testing.assert_allclose(out.float(), out2.float())
+
+n = 1
+dim1 = torch.randint(1,64, size=(n,)).tolist()
+dim2 = torch.randint(32,128, size=(n,)).tolist()
+dim3 = torch.randint(32,256, size=(n,)).tolist()
+values = list(product(dim1,dim2,dim3))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3", values, ids=names)
+def test_vector_quant(dim1, dim2, dim3):
+    dim2 = dim2 - (dim2 % 16)
+    dim3 = dim3 - (dim3 % 16)
+    for i in range(k):
+        A = torch.randn(size=(dim2, dim3), device='cuda')
+        qA, SA = F.vectorwise_quant(A, dim=0)
+        A1 = F.vectorwise_dequant(qA, SA)
+        torch.testing.assert_allclose(A1, A, atol=0.01, rtol=0.1)
+
+
+
+n = 2
+dim1 = torch.randint(2,256, size=(n,)).tolist()
+dim2 = torch.randint(2,256, size=(n,)).tolist()
+dim3 = torch.randint(2,256, size=(n,)).tolist()
+#dim1, dim2 = (256,), (256,)
+dtype = [torch.int8, torch.int32]
+a_order = ['row']
+out_order = ['col', 'row', 'col32']
+transpose = [False]
+dims = [2, 3]
+values = list(product(dim1,dim2,dim3, dims,dtype, a_order, out_order, transpose))
+
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dims_{3}_dtype_{4}_orderA_{5}_orderOut_{6}_transpose_{7}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dims, dtype, orderA, orderOut, transpose", values, ids=names)
+def test_nvidia_transform(dim1, dim2, dim3, dims, dtype, orderA, orderOut, transpose):
+    if dims == 3 and out_order != 'col32': return
+    if dtype == torch.int32 and out_order != 'col32': return
+    func = F.get_transform_func(dtype, orderA, orderOut, transpose)
+
+    if dims == 2:
+        A = torch.randint(-128, 127, size=(dim1, dim2), device='cuda').to(dtype)
+    elif dims == 3:
+        A = torch.randint(-128, 127, size=(dim1, dim2, dim3), device='cuda').to(dtype)
+
+    out, S = F.nvidia_transform(A, to_order=orderOut)
+
+    if orderOut == 'row':
+        torch.testing.assert_allclose(A.flatten(), out.flatten())
+    elif orderOut == 'col':
+        torch.testing.assert_allclose(A.t().flatten(), out.flatten())
+    elif orderOut == 'col32':
+        if dims == 2:
+            n = A.shape[0]*(A.shape[1] + (32 - (A.shape[1]%32)))
+        elif dims == 3:
+            n = A.shape[0]*A.shape[1]*(A.shape[2] + (32 - (A.shape[2]%32)))
+        assert out.numel() == n
+    elif orderOut == 'col_turing':
+        # 32 col 8 row tiles
+        n = (A.shape[0]+(8- A.shape[0]%8))*(A.shape[1] + (32 - (A.shape[1]%32)))
+        assert out.numel() == n
+        total_coltile = (A.shape[1] // 32) + (1 if A.shape[1] % 32 != 0 else 0)
+        for row in range(A.shape[0]):
+            for col in range(A.shape[1]):
+                i = row*A.shape[1]
+                j = col
+
+                coltile = (col // 32) + (1 if col % 32 != 0 else 0)
+                rowtile = ((row // 8) + (1 if row % 8 != 0 else 0))*total_coltile
+                offset = 32*8*(rowtile+coltile)
+                col2 = col % 32
+                row2 = (row%8)*32
+
+
+                assert A.flatten()[i+j] == A[row, col]
+                #assert A.flatten()[i+j] == out.flatten()[row2+col2]
+                #torch.testing.assert_allclose(A.flatten()[i+j], A[row, col])
+                #torch.testing.assert_allclose(A.flatten()[i+j], out.flatten()[row2+ col2+block_offset])
+
+    if orderOut == 'col32':
+        out2, S = F.nvidia_transform(out, from_order=orderOut, to_order='row', state=S)
+        torch.testing.assert_allclose(A, out2)
+
+
+n = 1
+dim1 = torch.randint(1,256, size=(n,)).tolist()
+dim2 = torch.randint(32,512, size=(n,)).tolist()
+dim3 = torch.randint(32,1024, size=(n,)).tolist()
+dim4 = torch.randint(32,1024, size=(n,)).tolist()
+
+#dim1 = [2]
+#dim2 = [2]
+#dim3 = [2]
+#dim4 = [2]
+
+dims = (2,3)
+ldb = [0]
+#ldb = list(range(256, 1*1024, 256))
+values = list(product(dim1,dim2,dim3,dim4,dims, ldb))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_dims_{4}_ldb_{5}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dim4, dims, ldb", values, ids=names)
+def test_igemmlt_int(dim1, dim2, dim3, dim4, dims, ldb):
+    for i in range(k):
+        if dims == 2:
+            A = torch.randint(-128, 127, size=(dim1, dim3), device='cuda').to(torch.int8)
+        elif dims == 3:
+            A = torch.randint(-128, 127, size=(dim1, dim2, dim3), device='cuda').to(torch.int8)
+        B = torch.randint(-128, 127, size=(dim4, dim3), device='cuda').to(torch.int8)
+        C1 = torch.matmul(A.float(), B.t().float())
+
+        A2, SA = F.transform(A, 'col32')
+        B2, SB = F.transform(B, 'col_turing')
+        C2, SC = F.igemmlt(A2, B2, SA, SB)
+        C3, S = F.nvidia_transform(C2, 'row', state=SC)
+        torch.testing.assert_allclose(C1, C3.float())
+
+        # transpose
+        B = torch.randint(-128, 127, size=(dim3, dim4), device='cuda').to(torch.int8)
+        C1 = torch.matmul(A.float(), B.float())
+
+        B2t, SBt = F.transform(B, 'col_turing', transpose=True)
+        C2, SC = F.igemmlt(A2, B2t, SA, SBt)
+        C3, S = F.nvidia_transform(C2, 'row', state=SC)
+        torch.testing.assert_allclose(C1, C3.float())
+
+dim1 = [32]
+dim2 = [32]
+dim3 = [32]
+dim4 = [32]
+
+dims = (2,)
+#ldb = list(range(256, 1*1024, 256))
+values = list(product(dim1,dim2,dim3,dim4,dims))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dim4_{3}_dims_{4}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dim4, dims", values, ids=names)
+def test_igemmlt_half(dim1, dim2, dim3, dim4, dims):
+    formatB = F.get_special_format_str()
+    for i in range(k):
+        if dims == 2:
+            A = torch.normal(0, 0.5, size=(dim1, dim3), device='cuda').half()
+        elif dims == 3:
+            A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device='cuda').half()
+        B = torch.randn((dim4, dim3), device='cuda').half()
+        torch.nn.init.xavier_uniform_(B)
+        C1 = torch.matmul(A, B.t())
+        C2 = bnb.matmul(A, B.t())
+
+        A = A.view(-1, A.shape[-1])
+
+        CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+        CB, CBt, statsB, statsBt, coo_tensor = F.double_quant(B)
+        C32A, SA = F.transform(CA, 'col32')
+        CxB, SB = F.transform(CB, to_order=formatB)
+        out1_32, Sout1_32 = F.igemmlt(C32A, CxB, SA, SB)
+        output = F.mm_dequant(out1_32, Sout1_32, statsAt, statsBt)
+
+        #print('')
+        #print(output.flatten()[:10])
+        #print(C1.flatten()[:10])
+        #print(C2.flatten()[:10])
+
+
+        #torch.testing.assert_allclose(C1.view(-1, C1.shape[-1]), output, atol=0.025, rtol=0.05)
+
+        # transpose
+        #B = torch.randint(-128, 127, size=(dim3, dim4), device='cuda').to(torch.int8)
+        #C1 = torch.matmul(A.float(), B.float())
+
+        #B2t, SBt = F.transform2(B, 'col_turing', transpose=True)
+        #C2, SC = F.igemmlt(A2, B2t, SA, SBt)
+        #C3, S = F.transform(C2, 'row', state=SC)
+        #torch.testing.assert_allclose(C1, C3.float())
+
+batch_size = 2
+seqdim = 512
+#values = [(batch_size, seqdim, 4*1024, 16*1024),(batch_size, seqdim, 5120, 4*5120),(batch_size, seqdim, 12*1024, 4*12*1024)]
+values = [(batch_size, seqdim, 4*1024, 3*4*1024),(batch_size, seqdim, 5120, 3*5120),(batch_size, seqdim, 12*1024, 4*12*1024)]
+
+
+#values = list(product(batch, seq, model, hidden))
+names = ['batch_{0}_seq_{1}_model_{2}_hidden_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("batch, seq, model, hidden", values, ids=names)
+def test_bench_8bit_training(batch, seq, model, hidden):
+    formatB = F.get_special_format_str()
+    A = torch.randn(batch, seq, model, device='cuda').half()
+    grad = torch.randn(batch, seq, model, device='cuda').half()
+    w1 = torch.randint(-128, 127, size=(hidden, model), device='cuda').half()
+    w2 = torch.randint(-128, 127, size=(model, hidden), device='cuda').half()
+    print('')
+
+    #torch.cuda.synchronize()
+    ## warmup
+    #for i in range(100):
+    #    torch.matmul(A, w1.t())
+    #torch.cuda.synchronize()
+
+    dtype = torch.int8
+    A = A.view(-1, A.shape[-1]).contiguous()
+    grad = grad.view(-1, grad.shape[-1]).contiguous()
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+
+        out1 = torch.matmul(A, w1.t()) # fc1
+        #out2 = torch.matmul(out1, w2.t())# fc2
+
+        #d1 = torch.matmul(grad, w2) # delta1
+        #d2 = torch.matmul(d1, w1) # delta2
+
+        #grad1 = torch.einsum('bo,bh->oh', out1, grad) # grad w2
+        #grad2 = torch.einsum('bh,bo->ho', A, d2) # grad w1
+
+    torch.cuda.synchronize()
+    t16 = time.time() - t0
+    print(t16)
+
+    #torch.cuda.empty_cache()
+
+    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
+    #Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)
+
+    #CTw1, Sw1 = F.transform2(Cw1, formatB)
+    #CTw2, Sw2 = F.transform2(Cw2, formatB)
+    #CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
+    #CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)
+
+    #CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+    #C32A, SA = F.transform2(CA, 'col32')
+    ## fc1
+    #out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype)
+    ##out1 = F.mm_dequant(out1_32, Sout1_32, statsAt, statsw1t)
+
+    ## fc2
+    #Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1)
+    #C32out1, Sout1 = F.transform2(Cout1, 'col32')
+    #out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype)
+    ##out2 = F.mm_dequant(out2_32, Sout2_32, statsout1t, statsw2t)
+
+    ## delta1
+    #Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad)
+    #C32grad, Sgrad = F.transform2(Cgrad, 'col32')
+    ##d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype)
+    ##d1 = F.mm_dequant(d1_32, Sd1_32, statsgradt, statsw2)
+
+    ## delta2
+    #Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1)
+    #C32d1, Sd1 = F.transform2(Cd1, 'col32')
+    ##d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype)
+    ##d2 = F.mm_dequant(d2_32, Sd2_32, statsd1t, statsw1)
+
+    ## grad1
+    #C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True)
+    #CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True)
+    ##grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype)
+    ##grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1, statsgrad)
+
+    ## grad2
+    #C32At, SAt = F.transform2(CAt, 'col32', transpose=True)
+    #CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True)
+    ##grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype)
+    ##grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsA, statsd1)
+
+    #Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)
+
+    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
+    #Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)
+
+    #CTw1, Sw1 = F.transform2(Cw1, formatB)
+    #CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)
+    #CTw2, Sw2 = F.transform2(Cw2, formatB)
+    #CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
+    #torch.cuda.synchronize()
+    #t0 = time.time()
+    #for i in range(k):
+    #    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
+    #    #CTw1, Sw1 = F.transform2(Cw1, formatB)
+    #    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
+    #    #CTw1, Sw1 = F.transform2(Cw1, formatB)
+
+    #    #CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A, threshold=3.5)
+    #    CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+    #    #CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)
+    #    #CTw2, Sw2 = F.transform2(Cw2, formatB)
+    #    #CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
+
+    #    C32A, SA = F.transform2(CA, 'col32')
+
+    #    # fc1
+    #    out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype)
+    #    #out1dn = F.mm_dequant(out1_32, Sout1_32, statsA, statsw1)
+
+    #    #print(coo_tensor.nnz)
+    #    #out1sp = F.spmm_coo(coo_tensor, w1.t())
+    #    #print(w1.t().shape)
+    #    #out1 = out1dn + out1sp
+
+    #    # fc2
+    #    Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1)
+    #    C32out1, Sout1 = F.transform2(Cout1, 'col32')
+    #    out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype)
+    #    #out2 = F.mm_dequant(out2_32, Sout2_32, statsout1, statsw2)
+
+    #    # delta1
+    #    Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad)
+    #    C32grad, Sgrad = F.transform2(Cgrad, 'col32')
+    #    d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype)
+    #    #d1 = F.mm_dequant(d1_32, Sd1_32, statsgrad, statsw2t)
+
+    #    # delta2
+    #    Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1)
+    #    C32d1, Sd1 = F.transform2(Cd1, 'col32')
+    #    d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype)
+    #    #d2 = F.mm_dequant(d2_32, Sd2_32, statsd1, statsw1t)
+
+    #    # grad1
+    #    #C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True)
+    #    #CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True)
+    #    #grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype)
+    #    #grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1t, statsgradt)
+
+    #    ## grad2
+    #    #C32At, SAt = F.transform2(CAt, 'col32', transpose=True)
+    #    #CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True)
+    #    #grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype)
+    #    #grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsAt, statsd1t)
+
+    #torch.cuda.synchronize()
+    #t8 = time.time() - t0
+    #print(t8)
+
+
+
+
+
+n = 2
+dim1 = torch.randint(64,256, size=(n,)).tolist()
+dim4 = torch.randint(64,1024, size=(n,)).tolist()
+
+#dim1 = [2*1024]
+#dim4 = [2*1024]
+
+#dim1 = [4]
+#dim4 = [4]
+
+dims = (2,)
+#ldb = list(range(256, 1*1024, 256))
+formatB = ['col_turing', 'col_ampere']
+values = list(product(dim1,dim4,dims, formatB))
+names = ['dim1_{0}_dim4_{1}_dims_{2}_formatB_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim4, dims, formatB", values, ids=names)
+def test_dequant_mm(dim1, dim4, dims, formatB):
+    inner = torch.randint(1, 128, size=(1,)).item()
+    formatB = F.get_special_format_str()
+    for i in range(k):
+        A = torch.randn(dim1, inner, device='cuda')
+        B = torch.randn(dim4, inner, device='cuda')
+        C1 = torch.matmul(A.half(), B.t().half())
+
+        A1, maxA = F.vectorwise_quant(A, dim=1)
+        B1, maxB = F.vectorwise_quant(B, dim=1)
+
+        A2, SA = F.nvidia_transform(A1, 'col32')
+        B2, SB = F.nvidia_transform(B1, formatB)
+        C2, SC = F.igemmlt(A2, B2, SA, SB)
+
+        C3, S = F.nvidia_transform(C2, 'row', state=SC)
+        C4 = F.vectorwise_mm_dequant(C3.float(), maxA, maxB.t())
+
+        count = (torch.isclose(C1, C4, atol=0.01, rtol=0.1) == 0).sum().item()
+        n = C1.numel()
+        p = 0.06
+        assert count/n < p, f'error in more than {p} of elements: {count}/{n}={count/n}'
+
+        C5 = F.mm_dequant(C2, SC, maxA.flatten(), maxB.flatten())
+        torch.testing.assert_allclose(C5, C4)
+        #print(C2)
+
+
+
+n = 2
+dim1 = [1*1024]
+dim2 = [1*1024]
+#dim1 = torch.randint(1,4*1024, size=(n,)).tolist()
+#dim2 = torch.randint(1,4*1024, size=(n,)).tolist()
+
+dims = (2,)
+#ldb = list(range(256, 1*1024, 256))
+values = list(product(dim1,dim2,dims))
+names = ['dim1_{0}_dim2_{1}_dims_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dims", values, ids=names)
+def test_colrow_absmax(dim1, dim2, dims):
+    for i in range(k):
+        threshold = 3.0
+        A = torch.randn(dim1, dim2, device='cuda').half()
+        A_truncated = A.clone()
+        A_truncated[torch.abs(A_truncated) >= 3.0] = 0.0
+        if dims == 2:
+            row_stats1, _ = torch.abs(A.float()).max(1)
+            col_stats1, _ = torch.abs(A.float()).max(0)
+            row_stats1_trunc, _ = torch.abs(A_truncated.float()).max(1)
+            col_stats1_trunc, _ = torch.abs(A_truncated.float()).max(0)
+        else:
+            assert False
+
+        row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=threshold)
+
+        A_blocked = einops.rearrange(torch.abs(A), '(rows row_tiles) (cols block_size)-> rows cols row_tiles block_size', row_tiles=16, block_size=64*4)
+        nnz_rows1_counts = (torch.abs(A_blocked)>=threshold).sum(3).flatten()
+        nnz_block_ptr1 = torch.zeros(nnz_rows1_counts.shape[0]+1, dtype=nnz_rows1_counts.dtype, device=nnz_rows1_counts.device)
+        nnz_block_ptr1[1:] = nnz_rows1_counts.cumsum(0)
+
+        torch.testing.assert_allclose(col_stats1_trunc, col_stats2)
+        torch.testing.assert_allclose(row_stats1_trunc, row_stats2)
+        torch.testing.assert_allclose(nnz_block_ptr1, nnz_block_ptr2)
+
+        row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=0.0)
+
+        torch.testing.assert_allclose(col_stats1, col_stats2)
+        torch.testing.assert_allclose(row_stats1, row_stats2)
+        assert nnz_block_ptr2 is None
+
+
+
+n = 2
+#dim1 = [8*1024]
+#dim2 = [4*1024]
+dim1 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim2 = torch.randint(1,4*1024, size=(n,)).tolist()
+
+values = list(product(dim1,dim2))
+names = ['dim1_{0}_dim2_{1}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2", values, ids=names)
+def test_double_quant(dim1, dim2):
+    for i in range(k):
+        A = torch.randn(dim1, dim2, device='cuda').half()
+        out_col1, Scol = F.vectorwise_quant(A, dim=0)
+        out_row1, Srow = F.vectorwise_quant(A, dim=1)
+
+        CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+
+        # max difference is 1 due to rounding differences
+        torch.testing.assert_allclose(CA, out_row1, atol=1, rtol=0)
+        torch.testing.assert_allclose(CAt, out_col1, atol=1, rtol=0)
+
+
+        n = CAt.numel()
+        num_not_close_rows = (torch.isclose(CA, out_row1, atol=1)==0).sum().item()
+        num_not_close_cols = (torch.isclose(CAt, out_col1, atol=1)==0).sum().item()
+
+        # allow for 1:500 error due to rounding differences
+        min_error = 1/500
+        if num_not_close_cols > (min_error*n):
+            print(f'Min error exceeded {num_not_close_cols} elements are different. Error: {num_not_close_cols/n:.4f}')
+            assert False
+        if num_not_close_rows > (min_error*n):
+            print(f'Min error exceeded {num_not_close_rows} elements are different. Error: {num_not_close_rows/n:.4f}')
+            assert False
+
+        torch.testing.assert_allclose(Srow.flatten(), statsA)
+        torch.testing.assert_allclose(Scol.flatten(), statsAt)
+
+
+n = 4
+dim1 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim4 = torch.randint(1,4*1024, size=(n,)).tolist()
+inner = torch.randint(1,4*1024, size=(n,)).tolist()
+
+dim1 = [6]
+dim4 = [4]
+inner = [8]
+
+values = list(zip(dim1, dim4, inner))
+names = ['dim1_{0}_dim4_{1}_inner_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim4, inner", values, ids=names)
+def test_integrated_igemmlt(dim1, dim4, inner):
+    for i in range(k):
+        A = torch.randn(dim1, inner, device='cuda').half()
+        B = torch.randn(dim4, inner, device='cuda').half()
+
+        out1 = torch.matmul(A.half(), B.t().half())
+
+        C1a, C1b, stats1a, stats1b, coo_tensor = F.double_quant(A)
+        C2a, C2b, stats2a, stats2b, coo_tensor = F.double_quant(B)
+        A1, maxA = F.vectorwise_quant(A, dim=1)
+        B1, maxB = F.vectorwise_quant(B, dim=1)
+
+        torch.testing.assert_allclose(maxA.flatten(), stats1a)
+        torch.testing.assert_allclose(maxB.flatten(), stats2a)
+        torch.testing.assert_allclose(C1a, A1, rtol=0, atol=1)
+        torch.testing.assert_allclose(C2a, B1, rtol=0, atol=1)
+
+        A2, SA = F.nvidia_transform(C1a, 'col32')
+        B2, SB = F.nvidia_transform(C2a, 'col_turing')
+        outC32, SC = F.igemmlt(A2, B2, SA, SB)
+        out2 = F.mm_dequant(outC32, SC, stats1a, stats2a)
+
+        A2, SA = F.nvidia_transform(A1, 'col32')
+        B2, SB = F.nvidia_transform(B1, 'col_turing')
+        C2, SC = F.igemmlt(A2, B2, SA, SB)
+
+        C3, S = F.nvidia_transform(C2, 'row', state=SC)
+        out3 = F.vectorwise_mm_dequant(C3.float(), maxA, maxB.t())
+
+        err1 = torch.abs(out1-out2).mean().item()
+        err2 = torch.abs(out1-out3).mean().item()
+        assert err2 <= err1*1.01
+
+
+n = 6
+dim1 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim4 = torch.randint(1,4*1024, size=(n,)).tolist()
+inner = torch.randint(1,4*1024, size=(n,)).tolist()
+
+values = list(zip(dim1, dim4, inner))
+names = ['dim1_{0}_dim4_{1}_inner_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim4, inner", values, ids=names)
+def test_igemmlt_row_scale(dim1, dim4, inner):
+    formatB = F.get_special_format_str()
+    err1, err2, err3 = [], [], []
+    relerr1, relerr2 = [], []
+    scale = 1
+    for i in range(k):
+        A = torch.randn(dim1, inner, device='cuda').half()
+        B = torch.randn(dim4, inner, device='cuda').half()
+        torch.nn.init.xavier_uniform_(B)
+        C1 = torch.matmul(A, B.t())
+
+        out1 = torch.matmul(A.half(), B.t().half())
+
+
+        C1a, C1b, stats1a, stats1b, coo_tensor = F.double_quant(A)
+        CB, absmaxB = F.vectorwise_quant(B, quant_type='linear')
+        A2, SA = F.nvidia_transform(C1a, 'col32')
+        B2, SB = F.nvidia_transform(CB, formatB)
+        A1, maxA = F.vectorwise_quant(A, dim=1)
+
+        c = 10.0*inner*scale
+        row_scale = torch.ones_like(maxA)/c
+        outC32, SC = F.igemmlt(A2, B2, SA, SB, dtype=torch.int8, row_scale=row_scale)
+        C3, S = F.nvidia_transform(outC32, 'row', state=SC)
+        maxval = torch.abs(C3).max()
+        if maxval == 127:
+            scale = 1.5
+        else:
+            scale = maxval/120
+        out3 = C3*maxA*absmaxB*c/(127*127)
+
+        C4 = torch.matmul(C1a.float(), CB.float().t())
+
+
+        C2a, C2b, stats2a, stats2b, coo_tensor = F.double_quant(B)
+        B2, SB = F.nvidia_transform(C2a, formatB)
+        outC32, SC = F.igemmlt(A2, B2, SA, SB)
+        out2 = F.mm_dequant(outC32, SC, stats1a, stats2a)
+
+        CA, SA = F.vectorwise_quant(A, dim=1, quant_type='vector')
+        CB, SB = F.vectorwise_quant(B, dim=1, quant_type='linear')
+
+        C = torch.matmul(CA.float(), CB.t().float())
+        out4 = C*SA*SB/(127*127)
+        #out4 = torch.clip(torch.round(C*SA/c), -127, 127)*c*SB/(127*127)
+
+        #print('='*80)
+        #print(out1)
+        #print(out2)
+        #print(out3)
+
+        #print(out1)
+        #print(out2)
+        #print(out3)
+        err1.append(torch.abs(out1-out2).mean().item())
+        err2.append(torch.abs(out1-out3).mean().item())
+        err3.append(torch.abs(out1-out4).mean().item())
+
+        #assert_all_approx_close(C3.float(), torch.round(C4*row_scale), rtol=0, atol=0, count=10)
+    print('')
+    print(sum(err1)/len(err1))
+    print(sum(err2)/len(err2))
+    print(sum(err3)/len(err3))
+
+
+dim1 = [1024, 2048]
+inner = [12288*4, 4096*4]
+dim4 = [12288, 4096]
+
+values = list(zip(dim1, dim4, inner))
+names = ['dim1_{0}_dim4_{1}_inner_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim4, inner", values, ids=names)
+def test_row_scale_bench(dim1, dim4, inner):
+    err1, err2, err3 = [], [], []
+    relerr1, relerr2 = [], []
+    scale = 1
+    A = torch.randn(dim1, inner, device='cuda').half()
+    B = torch.randn(dim4, inner, device='cuda').half()
+    torch.nn.init.xavier_uniform_(B)
+    # warmpup
+    for i in range(k):
+        C1 = torch.matmul(A, B.t())
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+        C1 = torch.matmul(A, B.t())
+    torch.cuda.synchronize()
+    print('16', time.time()-t0)
+
+    C1a, C1b, stats1a, stats1b, coo_tensor = F.double_quant(A)
+    CB, absmaxB = F.vectorwise_quant(B, quant_type='linear')
+    A2, SA = F.nvidia_transform(C1a, 'col32')
+    B2, SB = F.nvidia_transform(CB, formatB)
+    A1, maxA = F.vectorwise_quant(A, dim=1)
+
+    c = 10.0*inner*scale
+    row_scale = maxA/c
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+        outC32, SC = F.igemmlt(A2, B2, SA, SB, dtype=torch.int8, row_scale=row_scale)
+    torch.cuda.synchronize()
+    print('row-wise', time.time()-t0)
+
+
+    C2a, C2b, stats2a, stats2b, coo_tensor = F.double_quant(B)
+    B2, SB = F.nvidia_transform(C2a, formatB)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+        outC32, SC = F.igemmlt(A2, B2, SA, SB)
+    torch.cuda.synchronize()
+    print('vector-wise', time.time()-t0)
+
+
+
+
+n = 2
+dim1 = torch.randint(2,1024, size=(n,)).tolist()
+dim2 = torch.randint(2,1024, size=(n,)).tolist()
+#dim1 = [8*1024]
+#dim2 = [4*1024]
+
+dim3 = [0]
+dtype = [torch.int8]
+a_order = ['row']
+out_order = ['col32', 'col_turing', 'col_ampere']
+transpose = [False, True]
+dims = [2]
+values = list(product(dim1,dim2,dim3, dims,dtype, a_order, out_order, transpose))
+names = ['dim1_{0}_dim2_{1}_dim3_{2}_dims_{3}_dtype_{4}_orderA_{5}_orderOut_{6}_{7}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dim3, dims, dtype, orderA, orderOut, transpose", values, ids=names)
+def test_transform(dim1, dim2, dim3, dims, dtype, orderA, orderOut, transpose):
+    for i in range(k):
+        if dims == 2:
+            A = torch.randint(10, 99, size=(dim1, dim2), device='cuda').to(dtype)
+        elif dims == 3:
+            A = torch.randint(10, 99, size=(dim1, dim2, dim3), device='cuda').to(dtype)
+
+        A.view(-1)[-1] = -1
+        if transpose:
+            At = A.t().contiguous()
+            out1, S1 = F.nvidia_transform(At, to_order=orderOut)
+        else:
+            out1, S1 = F.nvidia_transform(A, to_order=orderOut)
+        out2, S2 = F.transform(A, to_order=orderOut, transpose=transpose)
+
+        assert S1[0][0] == S2[0][0]
+        assert S1[0][1] == S2[0][1]
+        #print(out1)
+        #print(out2)
+
+        torch.testing.assert_allclose(out1, out2)
+
+n = 2
+#dim1 = torch.randint(2,1024, size=(n,)).tolist()
+#dim2 = torch.randint(2,1024, size=(n,)).tolist()
+dim1 = [1]
+dim2 = [33]
+
+dtype = [torch.int8]
+#a_order = ['col_turing', 'col_ampere']
+a_order = ['col_turing']
+out_order = ['row']
+values = list(product(dim1,dim2,dtype, a_order, out_order))
+names = ['dim1_{0}_dim2_{1}_dtype_{2}_orderA_{3}_orderOut_{4}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dtype, orderA, orderOut", values, ids=names)
+def test_transform_to_row(dim1, dim2, dtype, orderA, orderOut):
+    for i in range(1):
+        A = torch.randint(-127, 127, size=(dim1, dim2), device='cuda').to(dtype)
+
+        out2, S2 = F.transform(A, to_order=orderA)
+        A2, S3 = F.transform(out2, from_order=orderA, to_order='row', state=S2)
+        assert A2.shape[0] == A.shape[0]
+        assert A2.shape[1] == A.shape[1]
+
+
+        print('')
+        print(A)
+        print(out2)
+        print(A2)
+
+
+        #torch.testing.assert_allclose(A, A2)
+
+
+
+
+def test_overflow():
+    formatB = F.get_special_format_str()
+    for i in range(2):
+        a = torch.arange(5, 15).cuda().to(torch.int8).view(-1,1 )
+        b = torch.arange(5, 15).cuda().to(torch.int8).view(-1,1 )
+
+        Ca, Sa = F.nvidia_transform(a, 'col32')
+        Cb, Sb = F.nvidia_transform(b, formatB)
+
+        c = F.igemmlt(Ca, Cb, Sa, Sb, dtype=torch.int8)
+        c2 = torch.matmul(a.float(), b.float().t())
+
+
+n = 2
+dim1 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim2 = torch.randint(1,4*1024, size=(n,)).tolist()
+#dim1 = [4]
+#dim2 = [5]
+
+values = list(product(dim1,dim2))
+names = ['dim1_{0}_dim2_{1}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2", values, ids=names)
+def test_coo_double_quant(dim1, dim2):
+    threshold = 3.00
+    for i in range(k):
+        A = torch.randn(dim1, dim2, device='cuda').half()
+
+        idx = (torch.abs(A) >= threshold)
+        CA2, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+        CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A, threshold=threshold)
+
+        if coo_tensor is not None:
+            A1 = A*idx
+            A2 = torch.zeros_like(A)
+            A2[coo_tensor.rowidx.long(), coo_tensor.colidx.long()] = coo_tensor.values
+            torch.testing.assert_allclose(A1, A2)
+
+            A1 = A*(idx==0)
+            A2 = (CA.float()*statsA.unsqueeze(1)/127).half()
+            torch.testing.assert_allclose(A*(idx==0), A2, rtol=0.05, atol=1.5e-2)
+
+n = 2
+dim1 = torch.randint(1,1*1024, size=(n,)).tolist()
+dim2 = torch.randint(1,1*1024, size=(n,)).tolist()
+#dim1 = [7]
+#dim2 = [11]
+transposed_B = [False, True]
+values = list(product(dim1,dim2, transposed_B))
+names = ['dim1_{0}_dim2_{1}_transposed_B_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, transposed_B", values, ids=names)
+def test_spmm_coo(dim1, dim2, transposed_B):
+    threshold = 1.5
+    dim3 = torch.randint(32, 128, size=(1,)).item()
+    #dim3 = 17
+    for i in range(k):
+        A = torch.randn(dim1, dim2).cuda().half()
+        if transposed_B:
+            B = torch.randn(dim3, dim2).cuda().half()
+        else:
+            B = torch.randn(dim2, dim3).cuda().half()
+
+        idx = torch.abs(A) >= threshold
+        nnz = (idx == 1).sum().item()
+        rows, cols = torch.where(idx)
+        values = A[idx]
+        cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
+        A2 = A*idx
+
+        if transposed_B:
+            out2 = F.spmm_coo(cooA, B.t())
+            out1 = torch.matmul(A2, B.t())
+        else:
+            out2 = F.spmm_coo(cooA, B)
+            out1 = torch.matmul(A2, B)
+
+        assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=30)
+
+
+
+def test_spmm_bench():
+    batch = 2
+    model = 1024*1
+    hidden = model*4
+    seq = 1024
+    dim1 = batch*seq
+    dim2 = model
+    dim3 = hidden
+    threshold = 4
+    A = torch.randn(dim1, dim2, device='cuda').half()
+    B = torch.randn(dim2, dim3, device='cuda').half()
     for i in range(10):
-        A1 = torch.randn(1024, 1024, device='cpu')
-        C, S = F.quantize_blockwise(A1)
-        A2 = F.dequantize_blockwise(C, S)
-        diff = torch.abs(A1-A2)
-        reldiff = diff/torch.abs(A1+1e-8)
-        diffs.append(diff.mean().item())
-        reldiffs.append(reldiff.mean().item())
-        assert diffs[-1] < 0.011
-    #print(sum(diffs)/len(diffs))
-    #print(sum(reldiffs)/len(reldiffs))
+        C1 = bnb.matmul(A, B)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+        C1 = bnb.matmul(A, B)
+    torch.cuda.synchronize()
+    t8 = time.time()-t0
+
+    idx = torch.abs(A) >= threshold
+    nnz = (idx == 1).sum().item()
+    print(nnz/idx.numel())
+    rows, cols = torch.where(idx)
+    values = A[idx]
+    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
 
-    diffs = []
     for i in range(10):
-        A1 = torch.rand(1024, 1024, device='cpu')
-        C, S = F.quantize_blockwise(A1)
-        A2 = F.dequantize_blockwise(C, S)
-        diff = torch.abs(A1-A2).mean().item()
-        assert diff < 0.0033
-        diffs.append(diff)
-        torch.testing.assert_allclose(A1, A2, atol=1e-2, rtol=0)
-    #print(sum(diffs)/len(diffs))
+        out2 = F.spmm_coo(cooA, B)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(k):
+        out2 = F.spmm_coo(cooA, B)
+    torch.cuda.synchronize()
+    tsp = time.time()-t0
+    print(tsp, t8)
+    print(tsp/t8)
+
+
+n = 2
+dim1 = torch.randint(256,1*1024, size=(n,)).tolist()
+dim2 = torch.randint(256,1*1024, size=(n,)).tolist()
+values = list(product(dim1,dim2))
+names = ['dim1_{0}_dim2_{1}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2", values, ids=names)
+def test_integrated_sparse_decomp(dim1, dim2):
+    threshold = 3.0
+    formatB = 'col_turing'
+    for i in range(k):
+        A = torch.randn(dim1, dim2).cuda().half()
+        w1 = torch.randn(dim1, dim2).cuda().half()
+        out1 = torch.matmul(A, w1.t())
+
+        Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
+        CTw1, Sw1 = F.transform(Cw1, formatB)
+
+        CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
+        C32A, SA = F.transform(CA, 'col32')
+
+        out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1)
+        out2 = F.mm_dequant(out1_32, Sout1_32, statsA, statsw1)
+
+        CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A, threshold=threshold)
+        C32A, SA = F.transform(CA, 'col32')
+
+        out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1)
+        out3 = F.mm_dequant(out1_32, Sout1_32, statsA, statsw1)
+
+        assert coo_tensor is not None
+
+        out4 = F.spmm_coo(coo_tensor, w1.t())
+        out5 = out3 + out4
+
+        err1 = torch.abs(out1-out2).mean().item()
+        err2 = torch.abs(out1-out5).mean().item()
+        assert err2 < err1
+
+
+def test_matmuls():
+    a = torch.randn(256, 256).half().cuda()
+    b = torch.randn(256, 256).half().cuda()
+    c1 = torch.matmul(a, b)
+    c2 = bnb.matmul(a, b)
+    c3 = bnb.matmul(a, b)
+
+    err1 = torch.abs(c1-c2).mean().item()
+    err2 = torch.abs(c1-c3).mean().item()
+    assert err1 < 0.2
+    assert err2 < 0.2
+
+
+
+n = 2
+#dim1 = torch.randint(1,1*1024, size=(n,)).tolist()
+#dim2 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim1 = [1*2048]
+dim2 = [12288]
+#dim1 = [32]
+#dim2 = [32]
+#dtype = [torch.float16, torch.int8]
+dtype = [torch.float16]
+out_function = ['zeros', 'ones']
+values = list(product(dim1,dim2, dtype, out_function))
+names = ['dim1_{0}_dim2_{1}_dtype_{2}_out_func_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dtype, out_func", values, ids=names)
+def test_spmm_coo_very_sparse(dim1, dim2, dtype, out_func):
+    out_func = getattr(torch, out_func)
+
+    threshold = 3.3
+    #threshold = 2.8
+    #threshold = 0.0
+    A = torch.randn(dim1, dim2, device='cuda').half()
+    if dtype == torch.float16:
+        B = torch.randn(dim2, dim2*4, device='cuda').half()
+        torch.nn.init.xavier_uniform_(B)
+    else:
+        B = torch.randn(dim2, dim2*4, device='cuda').half()
+        torch.nn.init.xavier_uniform_(B)
+        B, SB = F.vectorwise_quant(B, quant_type='linear')
+        #B = torch.randint(-127, 127, size=(dim2, dim2*4), device='cuda').to(torch.int8)
+
+    print('')
+    idx = torch.abs(A) >= threshold
+    nnz = (idx == 1).sum().item()
+    rows, cols = torch.where(idx)
+    values = A[idx]
+    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
+    A2 = A*idx
+    out1 = torch.matmul(A2.half(), B.half())
+    out = out_func(out1.shape, dtype=torch.float16, device=out1.device)
+    out1 += out.clone()
+    out2 = F.spmm_coo_very_sparse(cooA, B, out=out)
+    #print(B)
+    #print(out1)
+    #print(out2)
+    p = 200/(2048*12288*4)
+    n = out1.numel()
+    count = math.ceil(p*n)
+    std = out1.std()
+    out1 /= std
+    out2 /= std
+    assert_all_approx_close(out1, out2.half(), rtol=0.01, atol=3.0e-2, count=count)
+    #assert_all_approx_close(out1, out2.half(), rtol=0.05, atol=0.01, count=count)
+
+    idx_col = torch.randint(0, A2.shape[-1], size=(15,))
+
+    #torch.testing.assert_allclose(out1, out2.half(), rtol=0.05, atol=0.001)
+
+    #Bt = torch.randn(dim2*4, dim2, device='cuda').half()
+    #torch.cuda.synchronize()
+    #t0 = time.time()
+    #print(A2.shape, B.shape)
+    #for i in range(100):
+    #   #out3 = F.spmm_coo(cooA, Bt.t())
+    #   #out2 = F.spmm_coo(cooA, B)
+    #   #out2 = F.spmm_coo_very_sparse(cooA, B)
+    #   #out1 = torch.matmul(A, Bt.t())
+
+    #torch.cuda.synchronize()
+    #print(time.time() - t0)
+
+def test_layout():
+    a1 = torch.rand(16, 64, device='cuda', dtype=torch.float16)
+    a1 = torch.arange(16* 64, device='cuda').reshape(16, 64).byte()
+    a2, s2 = F.transform(a1, 'col_turing')
+    print(a2.shape)
+
+    print(a1.flatten()[8*64:8*64+32])
+    for i in range(4):
+        print(a2.flatten()[i*8*32:i*8*32+32], 0)
+
+
+def test_coo2csr():
+    threshold = 1
+    A = torch.randn(128, 128).half().cuda()
+    idx = torch.abs(A) >= threshold
+    nnz = (idx == 1).sum().item()
+    rows, cols = torch.where(idx)
+    values = A[idx]
+    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
+    A2 = A*idx
+    csrA = F.coo2csr(cooA)
+    counts = csrA.rowptr[1:] - csrA.rowptr[:-1]
+    assert counts.numel() == A.shape[0]
+
+    torch.testing.assert_allclose(counts, (A2!=0).sum(1))
+    idx = (A2!=0)
+    torch.testing.assert_allclose(A2[idx], csrA.values)
+
+
+def test_coo2csc():
+    threshold = 1
+    A = torch.randn(128, 128).half().cuda()
+    idx = torch.abs(A) >= threshold
+    nnz = (idx == 1).sum().item()
+    rows, cols = torch.where(idx)
+    values = A[idx]
+    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
+    A2 = A*idx
+    cscA = F.coo2csc(cooA)
+    counts = cscA.colptr[1:] - cscA.colptr[:-1]
+    assert counts.numel() == A.shape[1]
+
+    torch.testing.assert_allclose(counts, (A2!=0).sum(0))
+    # torch uses row-major -> use transpose to transfer to col-major
+    idx = (A2.t()!=0)
+    torch.testing.assert_allclose(A2.t()[idx], cscA.values)
+
+
+
+n = 2
+#dim1 = torch.randint(1,1*1024, size=(n,)).tolist()
+#dim2 = torch.randint(1,4*1024, size=(n,)).tolist()
+dim1 = [1*2048]
+#dim2 = [12288]
+dim2 = [2048]
+#dim1 = [2]
+#dim2 = [2]
+dtype = [torch.int8]
+values = list(product(dim1,dim2, dtype))
+names = ['dim1_{0}_dim2_{1}_dtype_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, dtype", values, ids=names)
+def test_spmm_coo_dequant(dim1, dim2, dtype):
+    threshold = 6.0
+    #threshold = 2.8
+    #threshold = 0.0
+    A = torch.randn(dim1, dim2, device='cuda').half()
+    B = torch.empty(dim2, dim2*4, device='cuda', dtype=torch.float16)
+    torch.nn.init.xavier_uniform_(B)
+    Bt = B.t().contiguous()
+
+
+    CB, CBt, statsB, statsBt, coo_tensor = F.double_quant(B)
+
+    rowidx = torch.randint(0, A.shape[-1], size=(15,))
+
+    A[:, rowidx] = 8.0
+
+    idx = torch.abs(A) >= threshold
+    nnz = (idx == 1).sum().item()
+    rows, cols = torch.where(idx)
+    values = A[idx]
+    cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values)
+    A2 = A*idx
+    out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
+    out1 = torch.matmul(A2, B.half())
+    out3 = F.spmm_coo_very_sparse(cooA, CBt.half())
+    out3 = out3*statsBt.half()/127
+
+    values, counts = torch.unique(cooA.rowidx, return_counts=True)
+    offset = counts.cumsum(0).int()
+    max_count, max_idx = torch.sort(counts, descending=True)
+    print(torch.median(max_count.float()))
+
+    torch.testing.assert_allclose(out2, out3, rtol=0.05, atol=0.001)
+
+    p = 200/(2048*12288*4)
+    n = out1.numel()
+    count = math.ceil(p*n)
+    assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=count)
+
+
+
+    #torch.cuda.synchronize()
+    #t0 = time.time()
+    #for i in range(100):
+    #   out2 = F.spmm_coo_very_sparse(cooA, B)
+    #torch.cuda.synchronize()
+    #print('fp16', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+       out2 = F.spmm_coo(cooA, B)
+    torch.cuda.synchronize()
+    print('cusparse fp16', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+       out2 = F.spmm_coo_very_sparse(cooA, CBt)
+    torch.cuda.synchronize()
+    print('int8', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+       out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
+    torch.cuda.synchronize()
+    print('int8+dequant', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+       out2 = torch.matmul(A, B)
+    torch.cuda.synchronize()
+    print('matmul', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        out1 = bnb.matmul(A, Bt)
+        out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt)
+        out = out1+out2
+    torch.cuda.synchronize()
+    print('sparse+ matmul', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        out1 = bnb.matmul(A, Bt)
+        torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1)
+    torch.cuda.synchronize()
+    print('partial matmul', time.time() - t0)
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        out1 = bnb.matmul(A, Bt)
+    torch.cuda.synchronize()
+    print('partial matmul', time.time() - t0)
+
+batch_size = 1
+seqdim = 2048
+values = []
+values.append((batch_size, seqdim, 768, 4*768))
+#values.append((batch_size, seqdim, 1024, 4*1024))
+#values.append((batch_size, seqdim, 1536, 4*1536))
+#values.append((batch_size, seqdim, 2048, 4*2048))
+#values.append((batch_size, seqdim, 2560, 4*2560))
+#values.append((batch_size, seqdim, 4096, 4*4096))
+#values.append((batch_size, seqdim, 5140, 4*5140))
+#values.append((batch_size, seqdim, 12288, 4*12288))
+names = ['batch_{0}_seq_{1}_model_{2}_hidden_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("batch, seq, model, hidden", values, ids=names)
+def test_bench_matmul(batch, seq, model, hidden):
+    formatB = F.get_special_format_str()
+
+    A = torch.randn(batch, seq, model, device='cuda').half()
+    B = torch.empty(hidden, model, dtype=torch.float16, device='cuda')
+    torch.nn.init.xavier_uniform_(B)
+
+    linear8bit = bnb.nn.Linear8bitLt(model, hidden, False).cuda().half()
+    linear8bit.eval()
+
+    outliers = torch.randint(0, model, size=(5,)).cuda()
+    A[:, :, outliers] = 8.0
+
+    linearMixedBit = bnb.nn.Linear8bitLt(model, hidden, False, threshold=6.0).cuda().half()
+    linearMixedBit.eval()
+
+    # warmup
+    for i in range(100):
+        torch.matmul(A, B.t())
+    torch.cuda.synchronize()
+    print('')
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        torch.matmul(A, B.t())
+    torch.cuda.synchronize()
+    print(f'pytorch: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        bnb.matmul(A, B)
+    torch.cuda.synchronize()
+    print(f'bnb lt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+    CA, CAt, SCA, SCAt, coo_tensorA = F.double_quant(A, threshold=0.0)
+    C32A, SA = F.transform(CA, 'col32')
+    CB, CBt, SCB, SCBt, coo_tensorB = F.double_quant(B)
+    CxB, SB = F.transform(CB, to_order=formatB)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
+    torch.cuda.synchronize()
+    print(f'igemmlt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+    BA, statsB = F.vectorwise_quant(B, dim=1)
+    CxB, SB = F.nvidia_transform(CB, to_order=formatB)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        A2 = A.view(-1, A.shape[-1]).contiguous()
+        CA, statsA = F.vectorwise_quant(A2, dim=1)
+        C32A, SA = F.nvidia_transform(CA, 'col32')
+        out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
+        Cout, Sout = F.nvidia_transform(out32, 'row', state=Sout32)
+        F.vectorwise_mm_dequant(Cout, statsA, statsB.t())
+    torch.cuda.synchronize()
+    print(f'vector pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+    BA, statsB = F.vectorwise_quant(B, dim=1, quant_type='linear')
+    CxB, SB = F.nvidia_transform(CB, to_order=formatB)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        A2 = A.view(-1, A.shape[-1]).contiguous()
+        CA, statsA = F.vectorwise_quant(A2, dim=1, quant_type='linear')
+        C32A, SA = F.nvidia_transform(CA, 'col32')
+        out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
+        Cout, Sout = F.nvidia_transform(out32, 'row', state=Sout32)
+        out = Cout*statsB*statsA*(1.0/(127*127))
+    torch.cuda.synchronize()
+    print(f'linear pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+    linear8bit(A)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        linear8bit(A)
+    torch.cuda.synchronize()
+    print(f'bnb linear8bitlt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+
+    linearMixedBit(A)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for i in range(100):
+        linearMixedBit(A)
+    torch.cuda.synchronize()
+    print(f'bnb linear8bitlt with threshold: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s')
+
+
+def test_zeropoint():
+    def min_max(x):
+        maxA = torch.amax(x, dim=1, keepdim=True)
+        minA = torch.amin(x, dim=1, keepdim=True)
+        midpoint = (maxA-minA)/2.0
+        dyna = 252/(maxA-minA)
+        #dyna *= 0.98
+        x = dyna*x
+        x = x - torch.round((dyna*(minA+midpoint)))
+        return x.to(torch.int8), minA, midpoint, dyna
+    batch = 2
+    seq = 2
+    model = 4
+    hidden = 2*model
+    #batch = 4
+    #seq = 2048
+    #model = 1024
+    #hidden = 8*model
+    A = torch.randn(batch*seq, model, device='cuda').half()-0.4
+    B = torch.nn.Parameter(torch.randn(model, hidden, device='cuda').half())
+
+    #A[0] = 0
+    #B[:, 0] = 0
+    #A = A*(A>0)
+    #A[0, 0] = 0
+    #A[0, 0] = 6.0
+
+    Ac, minA, midpoint, dyna = min_max(A)
+    #print(Ac[0, 0], 'zero')
+    #print(Ac, Ac.min(), Ac.max())
+    Bc, maxB = F.vectorwise_quant(B, quant_type='linear')
+    out = F.igemm(Ac, Bc)
+    out2 = torch.matmul(A,B)
+    offset = B.sum(0)*torch.round(dyna*(minA+midpoint))/dyna
+    out = out.float()
+    #print(out.shape, maxB.shape, scale.shape, offset.shape)
+    norm1 = maxB/127
+    C4 = (out/dyna)*norm1+offset
+
+
+    B1 = torch.nn.Parameter(B.clone())
+    B2 = torch.nn.Parameter(B.clone())
+    B3 = torch.nn.Parameter(B.clone())
+    B4 = torch.nn.Parameter(B.clone())
+
+
+    C1 = torch.matmul(A, B1)
+    C2 = bnb.matmul_cublas(A, B2, None, 'linear')
+    C3 = bnb.matmul_cublas(A, B3, None, 'zeropoint')
+    C4 = bnb.matmul_cublas(A, B4, None, 'vector-zeropoint')
+
+    err1 = torch.abs(C1-C2).mean().item()
+    err2 = torch.abs(C1-C3).mean().item()
+    err3 = torch.abs(C1-C4).mean().item()
+    print(err1, err2, err3)
+    #assert err1 > err2
+
+    loss1 = C1.mean()
+    loss2 = C2.mean()
+    loss3 = C3.mean()
+    loss4 = C4.mean()
+
+    loss1.backward()
+    loss2.backward()
+    loss3.backward()
+    loss4.backward()
+
+    print(B.grad)
+    print(B1.grad)
+    print(B2.grad)
+    print(B3.grad)
+    print(B4.grad)
+    err1 = torch.abs(B1.grad-B2.grad).mean().item()
+    err2 = torch.abs(B1.grad-B3.grad).mean().item()
+    err3 = torch.abs(B1.grad-B4.grad).mean().item()
+    print(err1, err2, err3)
+
+
+
+
+def test_zp():
+    def quant_zp(x):
+        dtype = x.dtype
+        x = x.float()
+        dyna = x.max() - x.min()
+        if dyna == 0: dyna = 1
+        qx = 254./dyna
+        minx = x.min()
+        #zpx = torch.round(minx* qx)
+        #zpx = 127 - torch.round(x.max()* qx)
+        zpx = torch.round(x.min()* qx) - 127
+        x = (qx*x) + zpx
+        return x, qx, zpx
+    batch = 2
+    seq = 512
+    model = 1024
+    hidden = 4*model
+    A = torch.randn(batch*seq, model, device='cuda').half()*0.1
+    B = torch.randn(model, hidden, device='cuda').half()*0.1
+
+
+    C0 = torch.matmul(A, B)
+
+
+    #A, SA = F.vectorwise_quant(A, quant_type='linear')
+    #B, SB = F.vectorwise_quant(B, quant_type='linear')
+    A = A.float()
+    B = B.float()
+
+    C1 = torch.matmul(A, B)
+    C3 = bnb.matmul(A.half(), B.t().contiguous().half())
+
+    zp = 1
+    #C2 = torch.matmul(A-zp, B)
+    #C2 += B.sum(0).view(1, -1)*zp
+    C2 = torch.matmul(A, B-zp)
+    C2 -= A.sum(1).view(-1, 1)*zp
+
+    ca, cqa, cza = quant_zp(A)
+    print(ca.min(), ca.max())
+    print((ca-cza).min(), (ca-cza).max())
+
+    zp = 1
+    scale = 2.0
+    C5 = torch.matmul((A*scale)-zp, B)
+    C5 += B.sum(0)*zp
+    C5 /= scale
+
+    CA, qa, zpa = quant_zp(A)
+    C4 = torch.matmul(CA, B)
+    C4 -= B.sum(0)*zpa
+    C4 /= qa
 
+    zpb = 1
+    zpa = 1
+    qa = 2
+    qb = 2
+    C6 = torch.matmul((A*qa)+zpa, (B*qb)+zpb)
+    C6 -= (qb*B.sum(0).view(1, -1)*zpa) + (qa*A.sum(1).view(-1, 1)*zpb)
+    C6 -= zpa*zpb*A.shape[1]
+    C6 /= qa*qb
 
-def test_histogram():
-    dim1, dim2 = 32, 32
-    source = torch.rand(dim1, dim2, device='cuda')
-    idx1 = torch.randint(0, 255, size=(dim1, dim2), device='cuda').int()
-    idx2 = torch.randint(0, 255, size=(dim1, dim2), device='cuda').int()
-    histogram1 = torch.zeros((256, 256)).cuda()
-    histogram2 = torch.zeros((256, 256)).cuda()
+    CA, qa, zpa = quant_zp(A)
+    CB, qb, zpb = quant_zp(B)
+    C7 = torch.matmul(CA, CB)
+    C7 -= (qb*B.sum(0).view(1, -1)*zpa) + (qa*A.sum(1).view(-1, 1)*zpb)
+    C7 -= zpa*zpb*A.shape[1]
+    C7 /= qa*qb
 
-    F.histogram_scatter_add_2d(histogram2, idx1, idx2, source)
+    print('')
+    #print(C0.flatten()[:10])
+    print(C1.flatten()[:10])
+    print(C2.flatten()[:10])
+    print(C3.flatten()[:10])
+    print(C5.flatten()[:10])
+    print(C6.flatten()[:10])
+    print(C7.flatten()[:10])
+    err1 = torch.abs(C1-C2).mean().item()
+    err2 = torch.abs(C1-C3).mean().item()
+    err3 = torch.abs(C1-C4).mean().item()
+    err4 = torch.abs(C1-C5).mean().item()
+    err5 = torch.abs(C1-C6).mean().item()
+    err6 = torch.abs(C1-C7).mean().item()
+    print(err1, err2, err3, err4, err5, err6)
 
-    for i in range(dim1):
-        for j in range(dim2):
-            histogram1[idx1[i, j].item(), idx2[i, j].item()] += source[i, j]
 
-    torch.testing.assert_allclose(histogram1, histogram2)
-    torch.testing.assert_allclose(histogram1.sum(), source.sum())
diff --git a/tests/test_modules.py b/tests/test_modules.py
index a0379cb..a2c950b 100644
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -1,42 +1,470 @@
-# Copyright (c) Facebook, Inc. and its affiliates. 
-#   
-# This source code is licensed under the MIT license found in the 
-# LICENSE file in the root directory of this source tree.
 import pytest
 import torch
+
+from itertools import product
+from torch import nn
+
 import bitsandbytes as bnb
 
+class MockArgs(object):
+    def __init__(self, initial_data):
+        for key in initial_data:
+            setattr(self, key, initial_data[key])
+
+class MLP8bit(torch.nn.Module):
+    def __init__(self, dim1, dim2, has_fp16_weights=True, threshold=0.0):
+        super(MLP8bit, self).__init__()
+        self.fc1 = bnb.nn.Linear8bitLt(dim1, dim2, has_fp16_weights=has_fp16_weights, threshold=threshold)
+        self.fc2 = bnb.nn.Linear8bitLt(dim2, dim1, has_fp16_weights=has_fp16_weights, threshold=threshold)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.fc2(x)
+        return x
+
+
+def get_args():
+    args = MockArgs([])
+    args.quant_type = 'vector'
+    args.use_8bit_training = 'full'
+    args.clip_freq = 9999
+    return args
+
+def assert_all_approx_close(a, b, atol=1e-8, rtol=1e-5, count=10):
+    idx = torch.isclose(a, b, rtol, atol)
+    sumval = (idx==0).sum().item()
+    if sumval > count:
+        print(f'Too many values not close: assert {sumval} < {count}')
+        torch.testing.assert_allclose(a, b, rtol, atol)
+
+class LinearFunction(torch.autograd.Function):
+
+    @staticmethod
+    def get_8bit_linear_trimmed(x, stochastic=False, trim_value=3.0):
+        round_func = LinearFunction.round_stoachastic if stochastic else torch.round
+        norm = math.sqrt(math.pi)/math.sqrt(2.0)
+        #std = torch.abs(x).mean()*norm
+        std = torch.std(x)
+        max1 = std*trim_value
+        x = x/max1*127
+        x = round_func(x)
+        x[x > 127] = 127
+        x[x < -127] = -127
+        x = x/127*max1
+
+        return x
+
+    def quant(x, quant_type, dim=1):
+        if quant_type == 'linear':
+            max1 = torch.abs(x).max().float()
+            xq = torch.round(x/max1*127).to(torch.int8)
+            return xq, max1
+        elif quant_type == 'vector':
+            max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True)
+            xq = torch.round(x/max1*127).to(torch.int8)
+            return xq, max1
+        elif quant_type == 'min-max':
+            maxA = torch.amax(x, dim=dim, keepdim=True).float()
+            minA = torch.amin(x, dim=dim, keepdim=True).float()
+            scale = (maxA-minA)/2.0
+            xq = torch.round(127*(x-minA-scale)/scale).to(torch.int8)
+            return xq, (minA.float(), scale.float())
+        else: return None
+
+    def dequant(xq, S1, S2, dtype, quant_type):
+        if quant_type == 'linear':
+            norm = S1*S2/(127*127)
+            # double cast needed to prevent overflows
+            return (xq.float()*norm).to(dtype)
+        elif quant_type == 'vector':
+            x = xq.float()
+            if len(xq.shape) == 2 and len(S1.shape) == 3: S1 = S1.squeeze(0)
+            if len(xq.shape) == 2 and len(S2.shape) == 3: S2 = S2.squeeze(0)
+            #print(x.shape, S1.shape, S2.shape)
+            if len(S1.shape) == 2:
+                x *= S1.t()/127
+            else:
+                x *= S1/127
+            x *= S2/127
+            return x.to(dtype)
+        else: return None
+
+    def dequant_min_max(xq, A, B, SA, SB, dtype):
+        offset = B.float().t().sum(0)*(SA[0]+SA[1])
+        x = xq.float()
+        if len(xq.shape) == 2 and len(SB.shape) == 3: SB = SB.squeeze(0)
+        if len(xq.shape) == 2 and len(SA.shape) == 3: SA = SA.squeeze(0)
+        if len(SB.shape) == 2:
+            x *= SB.t()/127
+        else:
+            x *= SB/127
+        x *= SA[1]/127
+        x +=offset
+        return x.to(dtype)
+
+
+    def get_8bit_linear(x, stochastic=False):
+        round_func = LinearFunction.round_stoachastic if stochastic else torch.round
+        max1 = torch.abs(x).max()
+        x = x/max1*127
+        x = round_func(x)/127*max1
+        #x = torch.round(x)/128*max1
+        return x
+
+    @staticmethod
+    def get_8bit_vector_wise(x, dim, stochastic=False):
+        round_func = LinearFunction.round_stoachastic if stochastic else torch.round
+        max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True)
+        max1[max1==0] = 1.0
+        x = (x*127)/max1
+        x = round_func(x)/127*max1
+        return x
+
+    @staticmethod
+    def round_stoachastic(x):
+        sign = torch.sign(x)
+        absx = torch.abs(x)
+        decimal = absx-torch.floor(absx)
+        rdm = torch.rand_like(decimal)
+        return sign*(torch.floor(absx)+(rdm < decimal).to(x.dtype))
+
+    @staticmethod
+    def fake_8bit_storage(w, exponent_bits):
+        code = bnb.functional.create_dynamic_map(n=exponent_bits).to(w.device)
+        absmax, C = bnb.functional.quantize_blockwise(w.data, code=code)
+        out = bnb.functional.dequantize_blockwise(absmax, C, code)
+        out = out.half()
+        w.copy_(out)
+        return out
+
+    @staticmethod
+    def fake_8bit_storage_quantile(w, args):
+        code = bnb.functional.estimate_quantiles(w.data, offset=args.offset)
+        #C = bnb.functional.quantize_no_absmax(code, w)
+        #out = bnb.functional.dequantize_no_absmax(code, C, out=w.data)
+        #print(out)
+        #out = out.half()
+        code /= torch.max(torch.abs(code))
+        absmax, C = bnb.functional.quantize_blockwise(w.data, code=code)
+        out = bnb.functional.dequantize_blockwise(absmax, C, code)
+        out = out.half()
+        w.copy_(out)
+        return out
+
+    @staticmethod
+    def fake_8bit_storage_stoachstic(w):
+        rand = torch.rand(1024, device=w.device)
+        absmax, C = bnb.functional.quantize_blockwise(w.data, rand=rand)
+        out = bnb.functional.dequantize_blockwise(absmax, C)
+        out = out.half()
+        w.copy_(out)
+        return out
+
+    @staticmethod
+    def fake_8bit_storage_with_max(w, topk=8):
+        blocked_w = einops.rearrange(w.flatten(), '(h b) -> h b', b=256)
+        max_val, idx = torch.sort(torch.abs(blocked_w), dim=1, descending=True)
+        idx = idx[:, :topk]
+        max_val = max_val[:, :topk]
+
+        mask = torch.zeros_like(blocked_w)
+        mask.scatter_(dim=1, index=idx, src=torch.ones_like(max_val))
+        mask = mask.bool()
+
+        # 1. zero out max values
+        # 2. quantize + dequantize
+        # 3. write back max values
+        # 4. copy matrix back to weight
+
+        values = blocked_w[mask]
+        blocked_w[mask] = 0
+
+        code = bnb.functional.create_dynamic_map()
+        code = code.to(w.device)
+        absmax, C = bnb.functional.quantize_blockwise(blocked_w.data)
+        bnb.functional.dequantize_blockwise(absmax, C, out=blocked_w)
+
+        blocked_w[mask] = values
+
+        unblocked_w = blocked_w.flatten().view(w.shape)
+
+        w.copy_(unblocked_w)
+        return unblocked_w
+
+
+    @staticmethod
+    def forward(ctx, x, weight, bias=None, args=None):
+        if args.use_8bit_training != 'off':
+            weight8, S1 = LinearFunction.quant(weight, args.quant_type, dim=1)
+            x8, S2 = LinearFunction.quant(x, args.quant_type, dim=2)
+            outputq = bnb.functional.igemm(x8, weight8.t())
+            output = LinearFunction.dequant(outputq, S1, S2, x.dtype, args.quant_type)
+            #if torch.rand(1) < 0.01:
+                #output32 = torch.matmul(x, weight.t())
+                #err = torch.abs(output-output32).float()
+                #relerr = err/(torch.abs(output32).float()+1e-8)
+                #print(f'{err.mean().item():.4f}, {relerr.mean().item():.4f}', args.quant_type, 'forward', proxy)
+        else:
+            #output = torch.matmul(x, weight.t())
+            output = torch.einsum('bsi,oi->bso', x, weight)
+
+        ctx.save_for_backward(x, weight, bias)
+        ctx.args = args
+
+        if bias is not None:
+            output += bias.unsqueeze(0).expand_as(output)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, weight, bias = ctx.saved_tensors
+        args = ctx.args
+        stochastic = False
+        grad_input = grad_weight = grad_bias = None
+        if bias is not None and ctx.needs_input_grad[2]: grad_bias = grad_output.sum(0)
+
+        # weight and x are already 8bit
+        # -> transform grad_output to 8-bit
+        if args.use_8bit_training == 'forward+wgrad':
+            grad_output8, S1 = LinearFunction.quant(grad_output, args.quant_type, dim=[0, 1])
+            x8, S2 = LinearFunction.quant(x, args.quant_type, dim=[0, 1])
+            grad_weight8 = bnb.functional.igemm(grad_output8, x8)
+            grad_weight = LinearFunction.dequant(grad_weight8, S1, S2, grad_output.dtype, args.quant_type)
+
+            #grad_weight32 = torch.einsum('bso,bsi->oi', grad_output, x)
+
+            grad_input = grad_output.matmul(weight)
+        elif args.use_8bit_training == 'full':
+            grad_output8, S1 = LinearFunction.quant(grad_output, args.quant_type, dim=[0, 1])
+            x8, S2 = LinearFunction.quant(x, args.quant_type, dim=[0, 1])
+            grad_weight8 = torch.zeros_like(weight, dtype=torch.int32)
+            bnb.functional.igemm(grad_output8, x8, out=grad_weight8)
+            grad_weight = LinearFunction.dequant(grad_weight8, S1, S2, grad_output.dtype, args.quant_type)
+
+            grad_output8, S1 = LinearFunction.quant(grad_output, args.quant_type, dim=2)
+            weight8, S3 = LinearFunction.quant(weight, args.quant_type, dim=0)
+            grad_input8 = bnb.functional.igemm(grad_output8, weight8)
+            grad_input = LinearFunction.dequant(grad_input8, S1, S3, grad_output.dtype, args.quant_type)
+
+        else:
+            grad_input = grad_output.matmul(weight)
+            grad_weight = torch.einsum('bsi,bso->oi', x, grad_output)
 
-@pytest.mark.parametrize("embcls", [bnb.nn.Embedding, bnb.nn.StableEmbedding], ids=['Embedding', 'StableEmbedding'])
-def test_embeddings(embcls):
-    bnb.optim.GlobalOptimManager.get_instance().initialize()
-    emb1 = torch.nn.Embedding(100, 512).cuda()
-    emb2 = embcls(100, 512).cuda()
+        return grad_input, grad_weight, grad_bias, None
 
-    adam1 = bnb.optim.Adam8bit(emb1.parameters())
-    adam2 = bnb.optim.Adam8bit(emb2.parameters())
+class Linear8bit(nn.Module):
+    def __init__(self, input_features, output_features, bias=True, args=None):
+        super(Linear8bit, self).__init__()
+        self.input_features = input_features
+        self.output_features = output_features
+        self.args = args
 
-    batches = torch.randint(1, 100, size=(100, 4, 32)).cuda()
+        self.weight = nn.Parameter(torch.empty(output_features, input_features))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(output_features))
+        else:
+            self.register_parameter('bias', None)
 
+        torch.nn.init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            torch.nn.init.zeros_(self.bias)
+
+    def forward(self, x):
+        self.args.training = self.training
+
+        return LinearFunction.apply(x, self.weight, self.bias, self.args)
+
+
+
+def test_linear8bit():
+    l0 = torch.nn.Linear(32, 64).cuda().half()
+    l1 = bnb.nn.Linear8bit(32,64, args=get_args()).cuda().half()
+    l2 = Linear8bit(32, 64, args=get_args()).cuda().half()
+    l3 = bnb.nn.Linear8bitLt(32,64).cuda().half()
+
+    l0.weight.data = l2.weight.data.clone()
+    l0.bias.data = l2.bias.data.clone()
+
+    l1.weight.data = l2.weight.data.clone()
+    l1.bias.data = l2.bias.data.clone()
+
+    l3.weight.data = l2.weight.data.clone()
+    l3.bias.data = l2.bias.data.clone()
+
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        t = torch.randn(16, 8, 64, device='cuda').half()
+        b2 = b1.clone()
+        b3 = b1.clone()
+        b0 = b1.clone()
+
+        o0 = l0(b0)
+        o1 = l1(b1)
+        o2 = l2(b2)
+        o3 = l3(b3)
+
+        assert_all_approx_close(o1, o2, atol=0.013, rtol=0.05, count=1)
+        assert_all_approx_close(o3, o2, atol=0.013, rtol=0.05, count=1)
+
+        loss0 = torch.nn.functional.mse_loss(o0, t)
+        loss1 = torch.nn.functional.mse_loss(o1, t)
+        loss2 = torch.nn.functional.mse_loss(o2, t)
+        loss3 = torch.nn.functional.mse_loss(o3, t)
+
+        loss0.backward()
+        loss1.backward()
+        loss2.backward()
+        loss3.backward()
+
+        assert_all_approx_close(l1.bias.grad, l2.bias.grad, atol=0.01, rtol=0, count=2)
+        assert_all_approx_close(l3.bias.grad, l2.bias.grad, atol=0.01, rtol=0, count=2)
+        assert_all_approx_close(l1.weight.grad, l2.weight.grad, atol=0.013, rtol=0.05, count=2)
+        assert_all_approx_close(l3.weight.grad, l2.weight.grad, atol=0.013, rtol=0.05, count=2)
+
+        err1 = torch.abs(l0.weight.grad-l1.weight.grad).mean().item()
+        err2 = torch.abs(l0.weight.grad-l2.weight.grad).mean().item()
+        err3 = torch.abs(l0.weight.grad-l3.weight.grad).mean().item()
+
+        assert err1*0.8 < err2
+        assert err2*0.8 < err3
+        assert err3*0.8 < err1
+
+        l0.weight.grad = None
+        l1.weight.grad = None
+        l2.weight.grad = None
+        l3.weight.grad = None
+        l0.bias.grad = None
+        l1.bias.grad = None
+        l2.bias.grad = None
+        l3.bias.grad = None
+
+
+threshold = [0.0, 3.0]
+values = threshold
+names = ['threshold_{0}'.format(vals) for vals in values]
+@pytest.mark.parametrize("threshold", values, ids=names)
+def test_linear8bitlt_inference(threshold):
+    l1 = bnb.nn.Linear8bitLt(32,64, threshold=threshold).cuda().half()
+    assert l1.weight.device.type == 'cuda'
+    assert l1.weight.dtype == torch.float16
+
+    l1.eval()
     for i in range(100):
-        batch = batches[i]
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = l1(b1)
+        if i == 1:
+            assert l1.state.CxB is not None
+
+def test_linear8bitlt_accumulated_gradient():
+    l1 = torch.nn.Sequential(*[bnb.nn.Linear8bitLt(32,32).cuda().half() for i in range(2)])
+    l2 = torch.nn.Sequential(*[torch.nn.Linear(32,32).cuda().half() for i in range(2)])
+    l2[0].weight = torch.nn.Parameter(l1[0].weight.clone())
+    l2[0].bias = torch.nn.Parameter(l1[0].bias.clone())
+    l2[1].weight = torch.nn.Parameter(l1[1].weight.clone())
+    l2[1].bias = torch.nn.Parameter(l1[1].bias.clone())
+    opt1 = bnb.optim.Adam8bit(l1.parameters(), lr=0.001)
+    opt2 = bnb.optim.Adam8bit(l2.parameters(), lr=0.001)
+
+    acc_steps = 10
+
 
-        embedded1 = emb1(batch)
-        embedded2 = emb2(batch)
+    for i in range(10):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = l1(b1)
+        o2 = l2(b1)
+        loss1 = o1.mean()
+        loss2 = o2.mean()
+        loss1.backward()
+        loss2.backward()
+        if i == 2:
+            assert l1[0].state.CxB is not None
+            assert l1[1].state.CxB is not None
 
-        l1 = embedded1.mean()
-        l2 = embedded2.mean()
+        if i > 0 and i % acc_steps == 0:
+            opt1.step()
+            opt1.zero_grad(True)
+            opt2.step()
+            opt2.zero_grad(True)
+            assert_all_approx_close(l1[0].weight, l2[0].weight, rtol=1.05, atol=0.01, count=2)
+            assert_all_approx_close(l1[1].weight, l2[1].weight, rtol=1.05, atol=0.01, count=2)
+            # we do this copy because otherwise we have small divergences over time that add up
+            l1[0].weight.data.copy_(l2[0].weight.data)
+            l1[1].weight.data.copy_(l2[1].weight.data)
+        else:
+            torch.testing.assert_allclose(l1[0].weight.grad, l2[0].weight.grad)
+            torch.testing.assert_allclose(l1[1].weight.grad, l2[1].weight.grad)
 
-        l1.backward()
-        l2.backward()
 
-        adam1.step()
-        adam2.step()
+threshold = [0.0, 2.0]
+values = threshold
+names = ['threshold_{0}'.format(vals) for vals in values]
+@pytest.mark.parametrize("threshold", values, ids=names)
+def test_linear8bitlt_no_fp16_weights(threshold):
+    l1 = bnb.nn.Linear8bitLt(32,64, threshold=threshold, has_fp16_weights=False).cuda().half()
+    assert l1.weight.dtype == torch.int8
 
-        adam1.zero_grad()
-        adam2.zero_grad()
+    l1.eval()
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = l1(b1)
+        assert o1.dtype == torch.float16
+
+    mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).cuda()
+    assert mlp.fc1.weight.dtype == torch.int8
+    assert mlp.fc2.weight.dtype == torch.int8
 
-        assert adam1.state[emb1.weight]['state1'].dtype == torch.uint8
-        assert adam2.state[emb2.weight]['state1'].dtype == torch.float32
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = mlp(b1)
+        assert o1.dtype == torch.float16
+        if threshold > 0: assert mlp.fc1.state.idx is not None
+        if threshold > 0: assert mlp.fc2.state.idx is not None
 
+    mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).cuda().half()
+    assert mlp.fc1.weight.dtype == torch.int8
+    assert mlp.fc2.weight.dtype == torch.int8
+
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = mlp(b1)
+        assert o1.dtype == torch.float16
+        if threshold > 0: assert mlp.fc1.state.idx is not None
+        if threshold > 0: assert mlp.fc2.state.idx is not None
 
+    mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).half().cuda()
+
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = mlp(b1)
+        assert o1.dtype == torch.float16
+        if threshold > 0: assert mlp.fc1.state.idx is not None
+        if threshold > 0: assert mlp.fc2.state.idx is not None
+    assert mlp.fc1.weight.dtype == torch.int8
+    assert mlp.fc2.weight.dtype == torch.int8
+
+
+    mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).half().to('cuda')
+
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = mlp(b1)
+        assert o1.dtype == torch.float16
+        if threshold > 0: assert mlp.fc1.state.idx is not None
+        if threshold > 0: assert mlp.fc2.state.idx is not None
+    assert mlp.fc1.weight.dtype == torch.int8
+    assert mlp.fc2.weight.dtype == torch.int8
+    assert mlp.fc1.weight.device.type == 'cuda'
+    assert mlp.fc2.weight.device.type == 'cuda'
+
+    mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).to(torch.float16).to('cuda')
+
+    for i in range(100):
+        b1 = torch.randn(16, 8, 32, device='cuda').half()
+        o1 = mlp(b1)
+        assert o1.dtype == torch.float16
+        if threshold > 0: assert mlp.fc1.state.idx is not None
+        if threshold > 0: assert mlp.fc2.state.idx is not None
+    assert mlp.fc1.weight.dtype == torch.int8
+    assert mlp.fc2.weight.dtype == torch.int8
+    assert mlp.fc1.weight.device.type == 'cuda'
+    assert mlp.fc2.weight.device.type == 'cuda'
diff --git a/tests/test_optim.py b/tests/test_optim.py
index c80fe51..b173eaa 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -1,12 +1,9 @@
-# Copyright (c) Facebook, Inc. and its affiliates. 
-#   
-# This source code is licensed under the MIT license found in the 
-# LICENSE file in the root directory of this source tree.
 import os
 import time
 import shutil
 import uuid
 import pytest
+import ctypes
 import torch
 import bitsandbytes as bnb
 import bitsandbytes.functional as F
@@ -14,7 +11,9 @@ import bitsandbytes.functional as F
 from os.path import join
 from itertools import product
 
-import apex
+#import apex
+
+k = 20
 
 def get_temp_dir():
     path = '/tmp/autoswap/{0}'.format(str(uuid.uuid4()))
@@ -26,55 +25,47 @@ def rm_path(path):
 
 str2optimizers = {}
 str2optimizers['adam_pytorch'] = (None, torch.optim.Adam, bnb.optim.Adam)
-str2optimizers['adam_apex'] = (None, apex.optimizers.FusedAdam, bnb.optim.Adam)
-str2optimizers['momentum_apex'] = (None, lambda pxx: apex.optimizers.FusedSGD(pxx, 0.01, 0.9), bnb.optim.Adam)
+#str2optimizers['adam_apex'] = (None, apex.optimizers.FusedAdam, bnb.optim.Adam)
+#str2optimizers['momentum_apex'] = (None, lambda pxx: apex.optimizers.FusedSGD(pxx, 0.01, 0.9), bnb.optim.Adam)
 str2optimizers['momentum_pytorch'] = (None, lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), bnb.optim.Adam)
-str2optimizers['lamb_apex'] = (None, lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.00, use_nvlamb=True), bnb.optim.Adam)
-str2optimizers['lars_apex'] = (None, lambda pxx: apex.parallel.LARC.LARC(apex.optimizers.FusedSGD(pxx, 0.01, 0.9)), bnb.optim.Adam)
+#str2optimizers['lamb_apex'] = (None, lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.00, use_nvlamb=True), bnb.optim.Adam)
+#str2optimizers['lars_apex'] = (None, lambda pxx: apex.parallel.LARC.LARC(apex.optimizers.FusedSGD(pxx, 0.01, 0.9)), bnb.optim.Adam)
 
 str2optimizers['adam'] = (torch.optim.Adam, bnb.optim.Adam)
-str2optimizers['adamw'] = (torch.optim.AdamW, bnb.optim.AdamW)
-str2optimizers['fused_adam'] = (apex.optimizers.FusedAdam, bnb.optim.Adam)
+#str2optimizers['fused_adam'] = (apex.optimizers.FusedAdam, bnb.optim.Adam)
 str2optimizers['momentum'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD(pxx, 0.01, 0.9, block_wise=False))
 str2optimizers['lars'] = (lambda pxx: bnb.optim.PytorchLARS(pxx, 0.01, 0.9), lambda pxx: bnb.optim.LARS(pxx, 0.01, 0.9))
-str2optimizers['lamb'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB)
+#str2optimizers['lamb'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB)
 str2optimizers['rmsprop'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop(pxx, 0.01, 0.9, block_wise=False))
-str2optimizers['adagrad'] = (lambda pxx: torch.optim.Adagrad(pxx, 0.01), lambda pxx: bnb.optim.Adagrad(pxx, 0.01, block_wise=False))
 str2optimizers['adam8bit'] = (torch.optim.Adam, lambda pxx: bnb.optim.Adam8bit(pxx, block_wise=False))
 str2optimizers['momentum8bit'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=False))
 str2optimizers['rmsprop8bit'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop8bit(pxx, 0.01, 0.9, block_wise=False))
-str2optimizers['lamb8bit'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB8bit)
+#str2optimizers['lamb8bit'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB8bit)
 str2optimizers['lars8bit'] = (lambda pxx: bnb.optim.PytorchLARS(pxx, 0.01, 0.9), lambda pxx: bnb.optim.LARS8bit(pxx, 0.01, 0.9))
 
 str2optimizers['adam8bit_blockwise'] = (torch.optim.Adam, lambda pxx: bnb.optim.Adam8bit(pxx, block_wise=True))
-str2optimizers['adamw8bit_blockwise'] = (torch.optim.Adam, lambda pxx: bnb.optim.AdamW8bit(pxx, block_wise=True))
 str2optimizers['momentum8bit_blockwise'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=True))
 str2optimizers['rmsprop8bit_blockwise'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop8bit(pxx, 0.01, 0.9, block_wise=True))
-str2optimizers['adagrad8bit_blockwise'] = (lambda pxx: torch.optim.Adagrad(pxx, 0.01), lambda pxx: bnb.optim.Adagrad8bit(pxx, 0.01, block_wise=True))
 
 str2statenames = {}
 str2statenames['adam'] = [('exp_avg', 'state1'), ('exp_avg_sq', 'state2')]
-str2statenames['adamw'] = [('exp_avg', 'state1'), ('exp_avg_sq', 'state2')]
 str2statenames['momentum'] = [('momentum_buffer', 'state1')]
 str2statenames['lars'] = [('momentum_buffer', 'state1')]
 str2statenames['lamb'] = [('exp_avg', 'state1'), ('exp_avg_sq', 'state2')]
 str2statenames['rmsprop'] = [('square_avg', 'state1')]
-str2statenames['adagrad'] = [('sum', 'state1')]
 str2statenames['adam8bit'] = [('exp_avg', 'state1', 'qmap1', 'max1'), ('exp_avg_sq', 'state2', 'qmap2', 'max2')]
 str2statenames['lamb8bit'] = [('exp_avg', 'state1', 'qmap1', 'max1'), ('exp_avg_sq', 'state2', 'qmap2', 'max2')]
 str2statenames['adam8bit_blockwise'] = [('exp_avg', 'state1', 'qmap1', 'absmax1'), ('exp_avg_sq', 'state2', 'qmap2', 'absmax2')]
-str2statenames['adamw8bit_blockwise'] = [('exp_avg', 'state1', 'qmap1', 'absmax1'), ('exp_avg_sq', 'state2', 'qmap2', 'absmax2')]
 str2statenames['momentum8bit'] = [('momentum_buffer', 'state1', 'qmap1', 'max1')]
 str2statenames['momentum8bit_blockwise'] = [('momentum_buffer', 'state1', 'qmap1', 'absmax1')]
 str2statenames['lars8bit'] = [('momentum_buffer', 'state1', 'qmap1', 'max1')]
 str2statenames['rmsprop8bit'] = [('square_avg', 'state1', 'qmap1', 'max1')]
 str2statenames['rmsprop8bit_blockwise'] = [('square_avg', 'state1', 'qmap1', 'absmax1')]
-str2statenames['adagrad8bit_blockwise'] = [('sum', 'state1', 'qmap1', 'absmax1')]
 
 dim1 = [1024]
 dim2 = [32, 1024, 4097, 1]
 gtype = [torch.float32, torch.float16]
-optimizer_names = ['adam', 'adamw', 'momentum', 'rmsprop', 'lars', 'lamb', 'adagrad']
+optimizer_names = ['adam', 'momentum', 'rmsprop', 'lars', 'lamb']
 values = list(product(dim1,dim2, gtype, optimizer_names))
 names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_{3}'.format(*vals) for vals in values]
 @pytest.mark.parametrize("dim1, dim2, gtype, optim_name", values, ids=names)
@@ -89,12 +80,12 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name):
     bnb_optimizer = str2optimizers[optim_name][1]([p2])
 
     if gtype == torch.float32:
-        atol, rtol = 2e-6, 1e-5
+        atol, rtol = 1e-6, 1e-5
     else:
         atol, rtol = 1e-4, 1e-3
 
 
-    for i in range(50):
+    for i in range(k):
         g = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.01
         p1.grad = g.clone().float()
         p2.grad = g.clone()
@@ -107,7 +98,7 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name):
 
         torch.testing.assert_allclose(p1, p2.float(), atol=atol, rtol=rtol)
 
-        if i % 10 == 0 and i > 0:
+        if i % (k//5) == 0 and i > 0:
             path = get_temp_dir()
             torch.save(bnb_optimizer.state_dict(),join(path, 'opt.pt'))
             del bnb_optimizer
@@ -148,7 +139,6 @@ def test_global_config(dim1, dim2, gtype):
     eps = 1e-8
 
     bnb.optim.GlobalOptimManager.get_instance().initialize()
-    bnb.optim.GlobalOptimManager.get_instance().override_config(p2, 'skip_zeros', True)
     bnb.optim.GlobalOptimManager.get_instance().override_config(p3, 'optim_bits', 8)
 
     bnb.optim.GlobalOptimManager.get_instance().register_parameters([p1, p2, p3])
@@ -163,8 +153,6 @@ def test_global_config(dim1, dim2, gtype):
     else:
         atol, rtol = 1e-4, 1e-3
 
-    original_p2 = p2[mask].clone()
-
     for i in range(50):
         g1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + 0.001
         g2 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + 0.001
@@ -173,38 +161,17 @@ def test_global_config(dim1, dim2, gtype):
         p2.grad = g2
         p3.grad = g3
 
-        if i > 30 and i % 10 == 0:
-            g1.data[mask] = 0.0
-            g2.data[mask] = 0.0
-            p1.grad = g1
-            p2.grad = g2
-            original_p1 = p1[mask].clone()
-            original_p2 = p2[mask].clone()
-            og_s1 = adam2.state[p2]['state1'][mask].clone()
-            og_s2 = adam2.state[p2]['state2'][mask].clone()
-            og_s11 = adam2.state[p1]['state1'][mask].clone()
-            og_s21 = adam2.state[p1]['state2'][mask].clone()
-
         adam2.step()
 
         assert adam2.state[p3]['state1'].dtype == torch.uint8
         assert adam2.state[p3]['state2'].dtype == torch.uint8
 
-        if i > 30 and i % 10 == 0:
-            torch.testing.assert_allclose(original_p2, p2[mask])
-            torch.testing.assert_allclose(adam2.state[p2]['state1'][mask], og_s1)
-            torch.testing.assert_allclose(adam2.state[p2]['state2'][mask], og_s2)
-            assert ((p1[mask]- original_p1)==0.0).sum() < p1.numel()
-            assert ((adam2.state[p1]['state1'][mask]- og_s11)==0.0).sum() == 0.0
-            assert ((adam2.state[p1]['state2'][mask]- og_s21)==0.0).sum() == 0.0
-
-
 
 
 dim1 = [1024]
 dim2 = [32, 1024, 4097]
 gtype = [torch.float32, torch.float16]
-optimizer_names = ['adam8bit', 'momentum8bit', 'rmsprop8bit', 'adam8bit_blockwise', 'adamw8bit_blockwise', 'lamb8bit', 'lars8bit', 'momentum8bit_blockwise', 'rmsprop8bit_blockwise', 'adagrad8bit_blockwise']
+optimizer_names = ['adam8bit', 'momentum8bit', 'rmsprop8bit', 'adam8bit_blockwise', 'lamb8bit', 'lars8bit', 'momentum8bit_blockwise', 'rmsprop8bit_blockwise']
 values = list(product(dim1,dim2, gtype, optimizer_names))
 names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_{3}'.format(*vals) for vals in values]
 @pytest.mark.parametrize("dim1, dim2, gtype, optim_name", values, ids=names)
@@ -370,13 +337,12 @@ def test_benchmark_blockwise(dim1, dim2, gtype, optim_name):
     if dim1 == 1 and dim2 == 1: return
     p1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1
 
-
     bnb_optimizer = str2optimizers[optim_name][1]([p1])
 
     g = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.01
     p1.grad = g
-    for i in range(5000):
-        if i == 500:
+    for i in range(k):
+        if i == k//5:
             # 100 iterations for burn-in
             torch.cuda.synchronize()
             t0 = time.time()
@@ -386,23 +352,8 @@ def test_benchmark_blockwise(dim1, dim2, gtype, optim_name):
     torch.cuda.synchronize()
     s = time.time()-t0
     print('')
-    params = 4500*4096*4096
+    params = (k-k//5)*dim1*dim2
     print(optim_name, gtype, s/params)
     #assert s < 3.9
 
 
-
-def test_str_betas():
-    betas = (0.80, 0.95)
-    strbetas = '(0.80, 0.95)'
-
-    layer = torch.nn.Linear(10, 10)
-
-    base = bnb.optim.Adam(layer.parameters(), betas=betas)
-    strbase = bnb.optim.Adam(layer.parameters(), betas=strbetas)
-    assert base.defaults['betas'][0] == 0.8
-    assert base.defaults['betas'][1] == 0.95
-    assert strbase.defaults['betas'][0] == 0.8
-    assert strbase.defaults['betas'][1] == 0.95
-
-