Function code1x16_dequant

inference_lib/src/aqlm/inference_kernels/cuda_kernel.cpp:184–227

}

torch::Tensor code1x16_dequant(
  const torch::Tensor& codes,
  const torch::Tensor& codebooks,
  const torch::Tensor& scales
) {
  check_use_bfloat16(codebooks);
  auto in_features = codes.size(1) * codebooks.size(3);
  auto out_features = scales.size(0);

  auto weight = torch::empty({out_features, in_features},
    torch::TensorOptions()
      .dtype(codebooks.dtype())
      .device(codebooks.device())
  );
  if (codebooks.size(3) == 8) {
    code1x16_dequant_cuda<8>(
      codes.data_ptr(),
      weight.data_ptr(),
      codebooks.data_ptr(),
      out_features,
      in_features
    );
  } else if (codebooks.size(3) == 16) {
    code1x16_dequant_cuda<16>(
      codes.data_ptr(),
      weight.data_ptr(),
      codebooks.data_ptr(),
      out_features,
      in_features
    );
  } else {
    throw c10::NotImplementedError(
      {__func__, __FILE__, static_cast<uint32_t>(__LINE__)},
      c10::str(
        "AQLM CUDA kernels only support codebooks with 8 or 16 features. Got ",
        codebooks.size(3),
        "."
      )
    );
  }
  weight *= scales.index({"...", 0, 0});

  return weight;
}

int4 accumulate_sizes(const torch::Tensor& codebook_partition_sizes)
{