new kernels

inference_lib/src/aqlm/cuda/__init__.py

@@ -1 +1 @@
-from .cuda_kernel import cuda_matmul
+from .cuda_kernel import cuda_gemm_1x16, cuda_gemm_2x8

inference_lib/src/aqlm/cuda/cuda_kernel.py

@@ -11,22 +11,25 @@
 )
 
 
-def cuda_gemm_stupid(
-    input: torch.Tensor,  #  [num_inputs, in_features]
+def cuda_gemm_1x16(
+    input: torch.Tensor,  #  [..., in_features]
     codes: torch.IntTensor,  #  [num_out_groups, num_in_groups, num_codebooks]
     codebooks: torch.Tensor,  #  [num_codebooks, codebook_size, out_group_size, in_group_size]
     scales: torch.Tensor,  #  [num_out_groups, 1, 1, 1]
     bias: Optional[torch.Tensor],
 ) -> torch.Tensor:
+    input_shape = input.shape
+    input = input.reshape(-1, input_shape[-1])
+
     device, dtype = codebooks.device, codebooks.dtype
     num_codebooks, codebook_size, out_group_size, in_group_size = codebooks.shape
     in_features = input.shape[1]
     out_features = codes.shape[0] * out_group_size
-    num_input_groups = codes.shape[1]
     assert input.ndim == 2
     assert scales.shape == (out_features // out_group_size, 1, 1, 1)
     assert in_features % in_group_size == 0
-    assert codebooks.shape[1] == 2**16
+    assert codebook_size == 2**16
+    assert num_codebooks == 1
     assert codes.dtype == torch.int16
     assert input.dtype == torch.float16 and codebooks.dtype == torch.float16
 
@@ -38,33 +41,37 @@ def cuda_gemm_stupid(
     output *= scales.flatten().unsqueeze(0)
     if bias is not None:
         output += bias
-    return output
+    return output.reshape(input_shape[:-1] + (-1,))
 
-    # codebook = torch.randn((codebook_size, in_group_size), dtype=torch.half, device=DEV)
-    # A = torch.randint(codebook_size, (out_features, in_features // in_group_size), dtype=torch.int, device=DEV)
-    # A_ref = torch.vstack([codebook[A[i]].flatten().unsqueeze(0) for i in range(M)])
-    # A = A.to(torch.int16)
-    # B = torch.randn((in_features, 1), dtype=torch.half, device=DEV)
-    # C = torch.zeros((out_features, 1), dtype=torch.half, device=DEV)
 
-    # C_ref = torch.matmul(A_ref, B)
-    # codebook_cuda.code16_matvec(A, B, C, codebook)
-
-
-def cuda_matmul(
-    input: torch.Tensor,
-    codes: torch.IntTensor,
-    codebooks: torch.Tensor,
-    scales: torch.Tensor,
+def cuda_gemm_2x8(
+    input: torch.Tensor,  #  [..., in_features]
+    codes: torch.IntTensor,  #  [num_out_groups, num_in_groups, num_codebooks]
+    codebooks: torch.Tensor,  #  [num_codebooks, codebook_size, out_group_size, in_group_size]
+    scales: torch.Tensor,  #  [num_out_groups, 1, 1, 1]
     bias: Optional[torch.Tensor],
 ) -> torch.Tensor:
     input_shape = input.shape
     input = input.reshape(-1, input_shape[-1])
 
-    return cuda_gemm_stupid(
-        input,
-        codes,
-        codebooks,
-        scales,
-        bias,
-    ).reshape(input_shape[:-1] + (-1,))
+    device, dtype = codebooks.device, codebooks.dtype
+    num_codebooks, codebook_size, out_group_size, in_group_size = codebooks.shape
+    in_features = input.shape[1]
+    out_features = codes.shape[0] * out_group_size
+    assert input.ndim == 2
+    assert scales.shape == (out_features // out_group_size, 1, 1, 1)
+    assert in_features % in_group_size == 0
+    assert codebook_size == 2**8
+    assert num_codebooks == 1
+    assert codes.dtype == torch.int8
+    assert input.dtype == torch.float16 and codebooks.dtype == torch.float16
+
+    output = torch.zeros(input.shape[0], out_features, device=device, dtype=dtype)
+    for i in range(input.shape[0]):
+        CUDA_KERNEL.code2x8_matvec(
+            codes.squeeze(2), input[i].unsqueeze(-1), output[i].unsqueeze(-1), codebooks.squeeze(0, 2)
+        )
+    output *= scales.flatten().unsqueeze(0)
+    if bias is not None:
+        output += bias
+    return output.reshape(input_shape[:-1] + (-1,))

inference_lib/src/aqlm/inference_kernels/kernel_selector.py

@@ -16,20 +16,24 @@ def forward_pass_quantized_linear(
     bias: Optional[torch.Tensor],
 ) -> torch.Tensor:
     num_codebooks, codebook_size, out_group_size, in_group_size = codebooks.shape
-    if cuda_kernel_applicable(input.is_cuda, num_codebooks, codebook_size, out_group_size, in_group_size):
-        from aqlm.cuda.cuda_kernel import cuda_matmul
-
-        return cuda_matmul(input, codes, codebooks, scales, bias)
-
-    if triton_kernel_applicable(input.is_cuda):
-        return triton_matmul(input, codes, codebooks, scales, bias)
-
-    dequantized_weight = _dequantize_weight(
-        unpack_int_data(codes, codebooks.shape[0].bit_length() - 1),
-        codebooks,
-        scales,
-    )
-    return F.linear(input, dequantized_weight, bias)
+    match (input.is_cuda, num_codebooks, codebook_size, out_group_size, in_group_size):
+        case (True, 1, 65536, 1, 8):
+            from aqlm.cuda.cuda_kernel import cuda_gemm_1x16
+
+            return cuda_gemm_1x16(input, codes, codebooks, scales, bias)
+        case (True, 2, 256, 1, 8):
+            from aqlm.cuda.cuda_kernel import cuda_gemm_2x8
+
+            return cuda_gemm_2x8(input, codes, codebooks, scales, bias)
+        case (True, _, _, _, _):
+            return triton_matmul(input, codes, codebooks, scales, bias)
+        case _:
+            dequantized_weight = _dequantize_weight(
+                unpack_int_data(codes, codebooks.shape[0].bit_length() - 1),
+                codebooks,
+                scales,
+            )
+            return F.linear(input, dequantized_weight, bias)
 
 
 def cuda_kernel_applicable(

main.py

@@ -8,7 +8,6 @@
 import torch.nn as nn
 from tqdm import trange
 from tqdm.auto import trange
-from transformers import PreTrainedModel
 
 from aq_engine import AQEngine
 from src.aq import QuantizedLinear
@@ -24,6 +23,7 @@
     get_sequential_groups,
 )
 from src.utils import using_tf32
+from transformers import PreTrainedModel
 
 try:
     import wandb