chore: import upstream snapshot with attribution

2026-07-13 12:35:23 +08:00
commit 1f0f055804
72 changed files with 147370 additions and 0 deletions
@@ -0,0 +1,99 @@
+import torch
+from torch.utils import benchmark
+from torch import nn
+
+from pack_weight import convert_weight_int8_to_int2
+from torch.profiler import profile, record_function, ProfilerActivity
+import ctypes
+import numpy as np
+# set all seed
+torch.manual_seed(42)
+np.random.seed(42)
+
+bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so')
+
+def bitnet_int8xint2_linear(input0, input1, s, ws, ret):
+    out_shape = list(input0.shape)
+    out_shape[-1] = input1.shape[0]
+
+    stream = torch.cuda.current_stream()
+
+    M = input0.shape[0]
+    if len(out_shape) == 3: 
+        M *= input0.shape[1]
+    N = input1.shape[0]
+    K = input1.shape[1] * 4
+
+    bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)])
+
+    return ret
+
+if __name__ == '__main__':
+    test_list = [
+        (2560,  2560), 
+        (3840,  2560), 
+        (13824, 2560),
+        (2560,  6912) ,
+        (3200, 3200), 
+        (4800, 3200), 
+        (3200, 10240),
+        (20480, 3200),
+    ]
+    for N,K in test_list:
+        weight = torch.randint(-1, 2, (N, K), dtype=torch.int8, device='cuda')
+        weight_scale = torch.ones(1, dtype=torch.bfloat16, device='cuda')
+        weight_compressed = convert_weight_int8_to_int2(weight).to('cuda')
+
+        for i in range(1):
+            input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
+            input0_bf16 = input0.to(torch.bfloat16)
+            input_np = input0.cpu().to(torch.int32).numpy()
+            weight_np = weight.cpu().to(torch.int32).T.numpy()
+            out_np = np.matmul(input_np,weight_np)
+            out_np = torch.tensor(out_np).cuda().to(torch.bfloat16)
+
+            s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
+            ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
+
+            ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
+            out = bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
+
+            print(f'custom == np {torch.all(out==out_np)}')
+
+        input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
+        input0_fp16 = input0.to(torch.float16)
+        input0_bf16 = input0.to(torch.bfloat16)
+        weight_fp16 = weight.to(torch.float16).T
+        weight_bf16 = weight.to(torch.bfloat16).T
+        ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
+        s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
+        ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
+        t0 = benchmark.Timer(
+            stmt="bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)",
+            setup="from __main__ import input0, weight_compressed, s, ws, ret, bitnet_int8xint2_linear",
+            num_threads=1,
+        )
+
+        t1 = benchmark.Timer(
+            stmt="torch.matmul(input0_bf16,weight_bf16)",
+            setup="from __main__ import input0_bf16, weight_bf16",
+            num_threads=1,
+        )
+
+        time0 = t0.timeit(50)
+        time1 = t1.timeit(50)
+
+        print(f'Shape{N,K}, W2A8: {time0.mean * 1e6:.2f}us, torch BF16: {time1.mean * 1e6:.2f}us')
+        # activities = [ ProfilerActivity.CUDA, 
+        #             #   ProfilerActivity.CPU
+        #               ]
+        # sort_by_keyword = 'cuda' + "_time_total"
+        # with profile(activities=activities, record_shapes=True) as prof:
+        #     with record_function("model_inference1"):
+        #         for _ in range(10):
+        #             bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
+        #             torch.matmul(input0_fp16,weight_fp16)
+        #             torch.matmul(input0_bf16,weight_bf16)
+
+        # print(prof.key_averages().table(sort_by=sort_by_keyword, row_limit=15))
+