| 4290 | # dynamic_decoder was set up with |
| 4291 | @cuda_stream_guard |
| 4292 | def decode(self, |
| 4293 | input_ids: torch.Tensor, |
| 4294 | context_lengths: torch.Tensor, |
| 4295 | sampling_config: SamplingConfig, |
| 4296 | prompt_embedding_table: torch.Tensor = None, |
| 4297 | tasks: torch.Tensor = None, |
| 4298 | prompt_vocab_size: torch.Tensor = None, |
| 4299 | stop_words_list=None, |
| 4300 | bad_words_list=None, |
| 4301 | streaming: bool = False, |
| 4302 | output_sequence_lengths: bool = False, |
| 4303 | output_generation_logits: bool = False, |
| 4304 | return_dict: bool = False, |
| 4305 | encoder_output: torch.Tensor = None, |
| 4306 | encoder_input_lengths: torch.Tensor = None, |
| 4307 | stopping_criteria: StoppingCriteria = None, |
| 4308 | logits_processor: LogitsProcessor = None, |
| 4309 | cross_attention_mask: List[torch.Tensor] = None, |
| 4310 | **kwargs): |
| 4311 | scfg = sampling_config |
| 4312 | batch_size = context_lengths.size(0) |
| 4313 | beam_width = scfg.num_beams |
| 4314 | max_context_length = torch.max(context_lengths).item() |
| 4315 | host_context_lengths = context_lengths.cpu() |
| 4316 | assert batch_size == self.batch_size, \ |
| 4317 | "Given batch size is different from the one used in setup()," \ |
| 4318 | "rerun the setup function with the new batch size to avoid buffer overflow." |
| 4319 | assert max_context_length <= self.max_context_length, \ |
| 4320 | "Given input length is large then the one used in setup()," \ |
| 4321 | "rerun the setup function with the new max_context_length to avoid buffer overflow." |
| 4322 | assert beam_width == self.beam_width, \ |
| 4323 | "Given beam width is different from the one used in setup()," \ |
| 4324 | "rerun the setup function with the new beam width to avoid buffer overflow." |
| 4325 | assert self.sink_token_length <= torch.min(context_lengths).item(), \ |
| 4326 | "Given sink token length is larger than shortest context length," \ |
| 4327 | "rerun the setup function with a smaller sink token length." |
| 4328 | ite = 0 # index of local batches, will always be 0 if pp_size = 1 |
| 4329 | |
| 4330 | if self.remove_input_padding and input_ids.dim() == 2: |
| 4331 | assert input_ids.shape[ |
| 4332 | 0] == 1, "Packed 2D input must have shape [1, <sum of input lengths>]" |
| 4333 | input_ids = input_ids.squeeze(0) |
| 4334 | |
| 4335 | self.__setup_decoder(input_ids, scfg, host_context_lengths) |
| 4336 | if not self.buffer_allocated: |
| 4337 | raise RuntimeError('Buffer not allocated, please call setup first!') |
| 4338 | |
| 4339 | sequence_limit_lengths = torch.full((batch_size, 1), |
| 4340 | self.max_seq_length, |
| 4341 | dtype=torch.int32, |
| 4342 | device=self.device) |
| 4343 | |
| 4344 | # Sequence_lengths for the dynamic decoder still has the input paddings. |
| 4345 | sequence_lengths = torch.full((batch_size * beam_width, 1), |
| 4346 | max_context_length, |
| 4347 | dtype=torch.int32, |
| 4348 | device=self.device) |
| 4349 | |