huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 11:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove unnecessary reinitialize to HeterogeneousNextTokenChooser to make sampling output correct

Browse Source 
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi A 2025-07-01 19:34:47 -07:00
parent 429dcd9c64
commit cf564ec0e2
1 changed files with 32 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
backends/gaudi/server/text_generation_server/models/flash_causal_lm.py
View File

@ -1076,6 +1076,7 @@ class FlashCausalLMBatch(Batch):
(0, padded_bs - self.cache_lengths_tensor.shape[0]), (0, padded_bs - self.cache_lengths_tensor.shape[0]),
value=0, value=0,
) )
if len(self.next_token_chooser.do_sample) != padded_bs:
next_token_chooser_parameters = [] next_token_chooser_parameters = []
next_token_chooser_parameters.extend([r.parameters for r in self.requests]) next_token_chooser_parameters.extend([r.parameters for r in self.requests])
pad_next_token_chooser_parameters(next_token_chooser_parameters, padded_bs) pad_next_token_chooser_parameters(next_token_chooser_parameters, padded_bs)
@ -1379,10 +1380,12 @@ class FlashCausalLMBatch(Batch):
self.all_input_ids_tensor[i] self.all_input_ids_tensor[i]
) )
self.all_input_ids_tensor = all_input_ids_tensor self.all_input_ids_tensor = all_input_ids_tensor
if len(self.next_token_chooser.do_sample) != max_padded_bs:
next_token_chooser_parameters = [] next_token_chooser_parameters = []
next_token_chooser_parameters.extend([r.parameters for r in self.requests]) next_token_chooser_parameters.extend([r.parameters for r in self.requests])
pad_next_token_chooser_parameters(next_token_chooser_parameters, max_padded_bs) pad_next_token_chooser_parameters(
next_token_chooser_parameters, max_padded_bs
)
# update past grammar states # update past grammar states
fsm_grammar_states = [0] * max_padded_bs fsm_grammar_states = [0] * max_padded_bs