chore: import upstream snapshot with attribution

gpu/pack_weight.py

@@ -0,0 +1,98 @@
+import torch
+import numpy as np
+
+
+def B_global_16x32_to_shared_load_16x32_layout(i, j):
+    """
+         stride * 8 * (tx // HALF_WARP_expr)
+                + (tx % 8) * stride
+                + 16 * ((tx % HALF_WARP_expr) // 8)
+    """
+    thread_id = i * 2 + j // 16
+    row = (thread_id // 16) * 8 + (thread_id % 8)
+    col = (j % 16) + 16 * ((thread_id % 16) // 8)
+    return row, col
+
+
+def permutate_weight_fastest(weight):
+    wmma_n = 16
+    wmma_k = 32
+    N = weight.shape[0]
+    K = weight.shape[1]
+    
+    # Create a lookup table for the permutation
+    mapping = np.zeros((wmma_n, wmma_k, 2), dtype=int)
+    for ii in range(wmma_n):
+        for jj in range(wmma_k):
+            mapping[ii, jj] = B_global_16x32_to_shared_load_16x32_layout(ii, jj)
+    
+    # Reshape weight for the final format
+    permutated_weight = np.zeros((N // wmma_n, K // wmma_k, wmma_n, wmma_k), dtype="int8")
+    
+    # Use advanced indexing for the entire operation
+    i_indices = np.arange(N // wmma_n)[:, np.newaxis, np.newaxis, np.newaxis]
+    j_indices = np.arange(K // wmma_k)[np.newaxis, :, np.newaxis, np.newaxis]
+    
+    # Create the source indices
+    src_i = i_indices * wmma_n + mapping[:, :, 0]
+    src_j = j_indices * wmma_k + mapping[:, :, 1]
+    
+    # Extract and reshape in one go
+    permutated_weight = weight[src_i, src_j]
+    
+    return permutated_weight
+
+
+def compress_int2_to_int8(int2_weight):
+    int8_weight = np.zeros(
+        (*int2_weight.shape[:-1], int2_weight.shape[-1] // 4), dtype=np.int8
+    )
+    for j in range(int2_weight.shape[-1] // 4):
+        for k in range(4):
+            int8_weight[:, :, :, j] |= int2_weight[:, :, :, j * 4 + k] << (k * 2)
+    return int8_weight
+
+
+def interleave_weight_int8(qweight, nbits=2):\
+    # reinterpret the data type of qweight to int32
+    # shift = [ 0,  8, 16, 24,  2, 10, 18, 26,  4, 12, 20, 28,  6, 14, 22, 30]
+    # index: [ 0,  4,  8, 12,  1,  5,  9, 13,  2,  6, 10, 14,  3,  7, 11, 15]
+    qweight = qweight.view(np.int32)
+    new_qweight = np.zeros_like(qweight)
+    bits_stride = 8
+    mask = (1 << nbits) - 1  # for 4bit the val is 0x0000000f
+    num_groups = 32 // bits_stride # 4
+    elems_per_group = bits_stride // nbits  # 4
+    for i in range(num_groups):
+        for j in range(elems_per_group):
+            offset = i * elems_per_group + j
+            shift = (offset % num_groups) * bits_stride + (offset // num_groups) * nbits
+
+            new_qweight |= ((qweight >> (nbits * offset)) & mask) << shift
+    return new_qweight.view(np.int8)
+
+
+
+def convert_weight_int8_to_int2(weight):
+    N = weight.shape[0]
+    K = weight.shape[1]
+
+    weight = weight+2
+    
+    weight = weight.cpu().numpy()
+
+    # print(weight)
+    # print(torch.max(weight), torch.min(weight))
+
+    # permutated_weight_slow = permutate_weight(weight)
+    permutated_weight = permutate_weight_fastest(weight)
+    # assert np.all(permutated_weight_slow == permutated_weight)
+    # print("Permutation is correct")
+    compressed_weight = compress_int2_to_int8(permutated_weight)
+    interleaved_weight = interleave_weight_int8(compressed_weight, 2)
+
+    ret = torch.from_numpy(interleaved_weight)
+
+    ret = torch.reshape(ret, (N, K // 4))
+
+    return ret