Class StaticBlockInferencePredictor

llm/predictor.py:1021–1166

class StaticBlockInferencePredictor(BlockInferencePredictorMixin, BasePredictor):
    def __init__(
        self,
        config: PredictorArgument,
        cache_kvs_shape: list[list[int]],
        tokenizer: PretrainedTokenizer = None,
    ):
        self.cache_kvs_shape = cache_kvs_shape
        BasePredictor.__init__(self, config, tokenizer)
        BlockInferencePredictorMixin.__init__(self, config, tokenizer)

        self.init_inputs(config)

        if config.export_precache:
            for i in range(self.num_layers):
                self.inputs["pre_caches_{}".format(i)] = self.pre_caches[i]

        self.cache_kvs = {}
        if config.use_cachekv_int8 == "dynamic" or config.use_cachekv_int8 == "static":
            for i in range(len(self.cache_kvs_shape) // 2):
                self.cache_kvs["key_caches_{}".format(i)] = paddle.zeros(self.cache_kvs_shape[2 * i], dtype="uint8")
                self.cache_kvs["value_caches_{}".format(i)] = paddle.zeros(
                    self.cache_kvs_shape[2 * i + 1], dtype="uint8"
                )
        else:
            for i in range(len(self.cache_kvs_shape) // 2):
                self.cache_kvs["key_caches_{}".format(i)] = paddle.zeros(
                    self.cache_kvs_shape[2 * i], dtype=config.dtype
                )
                self.cache_kvs["value_caches_{}".format(i)] = paddle.zeros(
                    self.cache_kvs_shape[2 * i + 1], dtype=config.dtype
                )

        for i in range(self.num_layers):
            if self.config.use_cachekv_int8 == "dynamic":
                self.inputs["k_quant_scales_" + str(i)] = self.k_quant_scales[i]
                self.inputs["v_quant_scales_" + str(i)] = self.v_quant_scales[i]
                self.inputs["k_dequant_scales_" + str(i)] = self.k_dequant_scales[i]
                self.inputs["v_dequant_scales_" + str(i)] = self.v_dequant_scales[i]

        self._create_predictor(config)
        self.input_names = self.predictor.get_input_names()

        self._share_data()
        self.seq_lens_handle = self.predictor.get_input_handle("seq_lens_this_time")

    def _create_predictor(self, predictor_args: PredictorArgument):
        if not is_paddlenlp_ops_available():
            raise ValueError(
                "you should install the paddlenlp ops to run inference predictor, "
                "https://github.com/PaddlePaddle/PaddleNLP/blob/develop/csrc/README.md"
            )

        infer_model_path = get_infer_model_path(predictor_args.model_name_or_path, predictor_args.model_prefix)

        config = paddle.inference.Config(infer_model_path + ".pdmodel", infer_model_path + ".pdiparams")

        config.switch_ir_optim(False)