remove unnecessary reinitialize to HeterogeneousNextTokenChooser to make sampling output correct

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

View File

@ -1076,6 +1076,7 @@ class FlashCausalLMBatch(Batch):
(0, padded_bs - self.cache_lengths_tensor.shape[0]),
value=0,
)
if len(self.next_token_chooser.do_sample) != padded_bs:
next_token_chooser_parameters = []
next_token_chooser_parameters.extend([r.parameters for r in self.requests])
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 = all_input_ids_tensor
if len(self.next_token_chooser.do_sample) != max_padded_bs:
next_token_chooser_parameters = []
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
fsm_grammar_states = [0] * max_padded_bs