Post-processing steps after completing a single token generation.
(
sampled_token_ids: paddle.Tensor,
model_output: ModelOutputData,
is_warmuping: bool,
is_chunk_step: paddle.Tensor,
enable_chunked_prefill: bool,
)
| 54 | |
| 55 | |
| 56 | def post_process_hpu( |
| 57 | sampled_token_ids: paddle.Tensor, |
| 58 | model_output: ModelOutputData, |
| 59 | is_warmuping: bool, |
| 60 | is_chunk_step: paddle.Tensor, |
| 61 | enable_chunked_prefill: bool, |
| 62 | ) -> None: |
| 63 | """Post-processing steps after completing a single token generation.""" |
| 64 | start_time = time.time() |
| 65 | |
| 66 | not_need_stop_hpu = model_output.not_need_stop.to(sampled_token_ids.place) |
| 67 | is_block_step_hpu = model_output.is_block_step.to(sampled_token_ids.place) |
| 68 | |
| 69 | update_inputs_v3( |
| 70 | model_output.stop_flags, |
| 71 | model_output.step_idx, |
| 72 | not_need_stop_hpu, |
| 73 | model_output.seq_lens_this_time, |
| 74 | model_output.seq_lens_encoder, |
| 75 | model_output.seq_lens_decoder, |
| 76 | model_output.max_dec_len, |
| 77 | model_output.input_ids, |
| 78 | model_output.stop_nums, |
| 79 | sampled_token_ids, |
| 80 | is_block_step_hpu, |
| 81 | model_output.eos_token_id, |
| 82 | model_output.next_tokens, |
| 83 | ) |
| 84 | |
| 85 | model_output.not_need_stop[:] = not_need_stop_hpu.cpu() |
| 86 | model_output.is_block_step[:] = is_block_step_hpu.cpu() |
| 87 | |
| 88 | end_time = time.time() |
| 89 | execution_time = (end_time - start_time) * 1000 |
| 90 | hpu_model_runner_profile_logger.info(f"post_process_hpu::update_inputs_v3 execution time(ms): {execution_time}") |
| 91 | |
| 92 | if is_warmuping: |
| 93 | return |
| 94 | start_time = time.time() |
| 95 | if enable_chunked_prefill: |
| 96 | sampled_token_ids = sampled_token_ids.cpu() |
| 97 | for i in range(is_chunk_step.shape[0]): |
| 98 | if is_chunk_step[i]: |
| 99 | sampled_token_ids[i] = -1 |
| 100 | save_output( |
| 101 | sampled_token_ids, |
| 102 | model_output.not_need_stop, |
| 103 | model_output.mp_rank, |
| 104 | ) |
| 105 | end_time = time.time() |
| 106 | execution_time = (end_time - start_time) * 1000 |
| 107 | hpu_model_runner_profile_logger.info(f"post_process_hpu::save_output execution time(ms): {execution_time}") |
| 108 | |
| 109 | |
| 110 | def recover_block_hpu( |
no test coverage detected