| 164 | self.bindings = self._allocate_memory(self.input_shapes, self.input_types, self.output_shapes, self.output_types) |
| 165 | |
| 166 | def forward(self, input_ids, *args, **kwargs): |
| 167 | bs = self.batch_size |
| 168 | max_length = self.max_sequence_length |
| 169 | TRTHFRunner.ENCODER_LENGTH = input_ids.shape[1] |
| 170 | input_length = input_ids.shape[1] |
| 171 | encoder_hidden_size = self.encoder_hidden_size |
| 172 | |
| 173 | # Check if the input data is on CPU (which usually means the PyTorch does not support current GPU). |
| 174 | is_cpu_mode = (input_ids.device == torch.device("cpu")) |
| 175 | |
| 176 | # We allocate the buffers using max_length, but we only need to first portion of it, so copy the data into the |
| 177 | # first portion of the input buffer. |
| 178 | # TODO: Could we just reuse input_ids' data_ptr() as the first binding when input_ids is already contiguous to |
| 179 | # avoid an additional D2D? |
| 180 | if is_cpu_mode: |
| 181 | self.inputs["input_ids"] = input_ids.int().flatten().contiguous().cuda() |
| 182 | self.bindings[0] = self.inputs["input_ids"].data_ptr() |
| 183 | else: |
| 184 | self.inputs["input_ids"][:bs * input_length] = input_ids.flatten() |
| 185 | |
| 186 | # Set the binding shape of input_ids, which should be (bs, input_length). |
| 187 | self.trt_context.set_binding_shape(0, input_ids.shape) |
| 188 | |
| 189 | # Launch TRT inference. |
| 190 | # TODO: Could we use execute_v2_async() instead of execute_v2()? |
| 191 | self.trt_context.execute_v2(bindings=self.bindings) |
| 192 | |
| 193 | # We allocate the buffers using max_length, but we only need to first portion of it, so get only the first |
| 194 | # portion of the output buffer and return that. |
| 195 | # TODO: Could we construct a Torch tensor using given data_ptr() to avoid this D2D copy? |
| 196 | hidden_states_output = self.outputs["hidden_states"] |
| 197 | if is_cpu_mode: |
| 198 | hidden_states_output = hidden_states_output.cpu() |
| 199 | |
| 200 | folded = hidden_states_output[:bs * input_length * encoder_hidden_size].view(bs, input_length, encoder_hidden_size) |
| 201 | |
| 202 | return folded |
| 203 | |
| 204 | class BARTTRTDecoder(TRTHFRunner): |
| 205 | |