Method forward

demo/HuggingFace/BART/trt.py:393–545

(self, input_ids, encoder_hidden_states, *args, **kwargs)

        return reordered_past

    def forward(self, input_ids, encoder_hidden_states, *args, **kwargs):
        # Get the batch size.
        bs = input_ids.shape[0] # in beam search mode, bs is batch_size * num_beams

        # Get the maximum sequence length.
        max_length = self.max_sequence_length

        # Get the vocab size.
        vocab_size = BARTModelTRTConfig.VOCAB_SIZE[self.variant]

        # Actual sequence length of the input_ids and the output hidden_states.
        input_length = input_ids.shape[1]

        # The sequence length of the encoder_hidden_states.
        encoder_length = TRTHFRunner.ENCODER_LENGTH

        # Encoder hidden size
        encoder_hidden_size = self.encoder_hidden_size

        # KV cache flag
        use_cache = kwargs.get("use_cache", False)

        # flag for switch between dual engines
        non_kv_flag = self.use_non_kv_engine or (self.config.use_cache and kwargs.get("past_key_values") is None)
        # condition 1: during e2e decoding test, based on flag
        # condition 2: during single-step decoder test, depending on whether past_key_values is empty
        # note: without --enable-kv-cache arg, this flag should remain False

        # denote as variable to allow switch between non-kv and kv engines in kv cache mode
        trt_context = self.trt_context_non_kv if non_kv_flag else self.trt_context
        bindings = self.bindings_non_kv if non_kv_flag else self.bindings
        inputs = self.inputs_non_kv if non_kv_flag else self.inputs
        outputs = self.outputs_non_kv if non_kv_flag else self.outputs

        # Check if the input data is on CPU (which usually means the PyTorch does not support current GPU).
        is_cpu_mode = (input_ids.device == torch.device("cpu")) or (self.return_device == "cpu")

        # We allocate the buffers using max_length, but we only need to first portion of it, so copy the data into the
        # first portion of the input buffer.
        # TODO: Could we just reuse input_ids' data_ptr() as the first binding when input_ids is already contiguous to
        # avoid an additional D2D?
        if is_cpu_mode:
            inputs["input_ids"] = input_ids.int().flatten().contiguous().cuda()
            bindings[0] = inputs["input_ids"].data_ptr()
        else:
            inputs["input_ids"][:bs * input_length] = input_ids.flatten()

        # Set the binding shape of input_ids, which should be (bs, input_length).
        trt_context.set_binding_shape(0, input_ids.shape)

        # If encoder hidden states have not been copied yet, copy the hidden states to the input buffer.
        if not self.persist_encoder_hidden_states:
            if is_cpu_mode:
                inputs["encoder_hidden_states"] = encoder_hidden_states.flatten().contiguous().cuda()
                bindings[1] = inputs["encoder_hidden_states"].data_ptr()
            else:
                inputs["encoder_hidden_states"][:bs * encoder_length * encoder_hidden_size] = encoder_hidden_states.flatten()