2x8 matmat cpp

inference_lib/src/aqlm/inference_kernels/cuda_kernel.cpp

@@ -88,8 +88,42 @@ void code2x8_matvec(
   );
 }
 
+torch::Tensor code2x8_matmat(
+  const torch::Tensor& input,
+  const torch::Tensor& codes,
+  const torch::Tensor& codebooks,
+  const torch::Tensor& scales
+) {
+  auto input_sizes = input.sizes();
+  auto out_features = codes.size(0) * codebooks.size(2);
+  auto flat_input = input.reshape({-1, input.size(-1)});
+  auto flat_output = torch::empty({flat_input.size(0), out_features},
+    torch::TensorOptions()
+      .dtype(input.dtype())
+      .device(input.device())
+  );
+
+  for (int i = 0; i < flat_input.size(0); ++i) {
+    auto input_vec = flat_input.index({i});
+    auto output_vec = flat_output.index({i});
+    code2x8_matvec(
+      codes.squeeze(2),
+      input_vec,
+      output_vec,
+      codebooks
+    );
+  }
+  flat_output *= scales.flatten().unsqueeze(0);
+  auto output_sizes = input_sizes.vec();
+  output_sizes.pop_back();
+  output_sizes.push_back(-1);
+  auto output = flat_output.view(output_sizes);
+  return output;
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("code1x16_matvec", &code1x16_matvec, "1x16 (2bit) codebook matrix-vector product.");
   m.def("code1x16_matmat", &code1x16_matmat, "1x16 (2bit) codebook matrix-matrix product.");
   m.def("code2x8_matvec", &code2x8_matvec, "2x8 (2bit) codebook matrix-vector product.");
+  m.def("code2x8_matmat", &code2x8_matmat, "2x8 (2bit) codebook matrix-matrix product.");
 }

inference_lib/src/aqlm/inference_kernels/kernel_selector.py

@@ -20,9 +20,9 @@ def forward_pass_quantized_linear(
 
             return CUDA_KERNEL.code1x16_matmat(input, codes, codebooks, scales) + (bias if bias is not None else 0)
         case (True, 2, 256, 1, 8):
-            from .cuda_kernel import cuda_gemm_2x8
-
-            return cuda_gemm_2x8(input, codes, codebooks, scales, bias)
+            from .cuda_kernel import CUDA_KERNEL
+            
+            return CUDA_KERNEL.code2x8_matmat(input, codes, codebooks, scales) + (bias if bias is not None else 0)
         case (True, _, _, _, _):
             from .triton_kernel import triton_matmul