(self,
*,
batch_size: int,
scfg: SamplingConfig,
sequence_lengths: torch.Tensor,
context_lengths: torch.Tensor,
host_context_lengths,
max_context_length: int,
beam_width: int,
cache_indirections: list,
input_ids: torch.Tensor,
hidden_states: torch.Tensor,
prompt_embedding_table: torch.Tensor,
tasks: torch.Tensor,
prompt_vocab_size: torch.Tensor,
ite: int,
sequence_limit_lengths: torch.Tensor,
stop_words_data,
bad_words_data,
output_sequence_lengths: bool = False,
output_generation_logits: bool = False,
return_dict: bool = False,
encoder_output: torch.Tensor = None,
encoder_input_lengths: torch.Tensor = None,
stopping_criteria: StoppingCriteria = None,
logits_processor: LogitsProcessor = None,
cross_attention_mask: List[torch.Tensor] = None,
**kwargs)
| 3966 | fmt=txt_format) |
| 3967 | |
| 3968 | def decode_regular(self, |
| 3969 | *, |
| 3970 | batch_size: int, |
| 3971 | scfg: SamplingConfig, |
| 3972 | sequence_lengths: torch.Tensor, |
| 3973 | context_lengths: torch.Tensor, |
| 3974 | host_context_lengths, |
| 3975 | max_context_length: int, |
| 3976 | beam_width: int, |
| 3977 | cache_indirections: list, |
| 3978 | input_ids: torch.Tensor, |
| 3979 | hidden_states: torch.Tensor, |
| 3980 | prompt_embedding_table: torch.Tensor, |
| 3981 | tasks: torch.Tensor, |
| 3982 | prompt_vocab_size: torch.Tensor, |
| 3983 | ite: int, |
| 3984 | sequence_limit_lengths: torch.Tensor, |
| 3985 | stop_words_data, |
| 3986 | bad_words_data, |
| 3987 | output_sequence_lengths: bool = False, |
| 3988 | output_generation_logits: bool = False, |
| 3989 | return_dict: bool = False, |
| 3990 | encoder_output: torch.Tensor = None, |
| 3991 | encoder_input_lengths: torch.Tensor = None, |
| 3992 | stopping_criteria: StoppingCriteria = None, |
| 3993 | logits_processor: LogitsProcessor = None, |
| 3994 | cross_attention_mask: List[torch.Tensor] = None, |
| 3995 | **kwargs): |
| 3996 | kv_cache_block_offsets = None |
| 3997 | host_kv_cache_block_offsets = None |
| 3998 | cross_kv_cache_block_offsets = None |
| 3999 | host_cross_kv_cache_block_offsets = None |
| 4000 | attention_mask = None |
| 4001 | outputs_context_logits = None |
| 4002 | outputs_generation_logits = [] |
| 4003 | |
| 4004 | def get_outputs_dict(output_ids, num_steps=self.max_new_tokens): |
| 4005 | outputs = {} |
| 4006 | outputs['output_ids'] = output_ids |
| 4007 | if scfg.output_log_probs: |
| 4008 | outputs['log_probs'] = self.log_probs |
| 4009 | if scfg.output_cum_log_probs: |
| 4010 | outputs['cum_log_probs'] = self.cum_log_probs |
| 4011 | if output_sequence_lengths: |
| 4012 | outputs[ |
| 4013 | 'sequence_lengths'] = self.sequence_length_buffer.reshape( |
| 4014 | [batch_size, beam_width]) |
| 4015 | if self.gather_context_logits: |
| 4016 | outputs['context_logits'] = outputs_context_logits |
| 4017 | if self.gather_generation_logits or output_generation_logits: |
| 4018 | outputs['generation_logits'] = outputs_generation_logits |
| 4019 | if self.is_medusa_mode or self.is_redrafter_mode: |
| 4020 | outputs['steps_to_finish'] = num_steps |
| 4021 | if self.is_medusa_mode: |
| 4022 | outputs['medusa_output_tokens'] = self.medusa_output_tokens |
| 4023 | outputs['accept_lengths'] = self.accept_lengths |
| 4024 | if self.medusa_temperature != 0.0: |
| 4025 | outputs['medusa_output_logits'] = self.medusa_output_logits |
no test coverage detected