Post-processing steps after completing a single token generation.
(
sampler_or_pooler_output: Union[SamplerOutput, PoolerOutput],
model_output: ModelOutputData,
share_inputs: InputBatch,
sampling_metadata: SamplingMetadata = None,
block_size: int = 64,
save_each_rank: bool = False,
speculative_decoding: bool = False,
skip_save_output: bool = False,
async_output_queue: queue.Queue = None,
think_end_id: int = -1,
splitwise_role_is_decode: bool = False,
enable_entropy: bool = False,
routing_replay_manager: RoutingReplayManager = None,
sampling_mask_zmq_client: ZmqIpcClient = None,
)
| 585 | |
| 586 | |
| 587 | def post_process( |
| 588 | sampler_or_pooler_output: Union[SamplerOutput, PoolerOutput], |
| 589 | model_output: ModelOutputData, |
| 590 | share_inputs: InputBatch, |
| 591 | sampling_metadata: SamplingMetadata = None, |
| 592 | block_size: int = 64, |
| 593 | save_each_rank: bool = False, |
| 594 | speculative_decoding: bool = False, |
| 595 | skip_save_output: bool = False, |
| 596 | async_output_queue: queue.Queue = None, |
| 597 | think_end_id: int = -1, |
| 598 | splitwise_role_is_decode: bool = False, |
| 599 | enable_entropy: bool = False, |
| 600 | routing_replay_manager: RoutingReplayManager = None, |
| 601 | sampling_mask_zmq_client: ZmqIpcClient = None, |
| 602 | ) -> None: |
| 603 | """Post-processing steps after completing a single token generation.""" |
| 604 | |
| 605 | if isinstance(sampler_or_pooler_output, PoolerOutput): |
| 606 | post_process_pooling( |
| 607 | sampler_or_pooler_output, |
| 608 | model_output, |
| 609 | share_inputs, |
| 610 | block_size, |
| 611 | save_each_rank, |
| 612 | skip_save_output, |
| 613 | async_output_queue, |
| 614 | routing_replay_manager, |
| 615 | ) |
| 616 | else: |
| 617 | if speculative_decoding: |
| 618 | post_process_specualate( |
| 619 | sampler_or_pooler_output, |
| 620 | model_output, |
| 621 | share_inputs, |
| 622 | sampling_metadata, |
| 623 | save_each_rank, |
| 624 | skip_save_output, |
| 625 | think_end_id, |
| 626 | splitwise_role_is_decode, |
| 627 | enable_entropy, |
| 628 | routing_replay_manager, |
| 629 | sampling_mask_zmq_client, |
| 630 | ) |
| 631 | else: |
| 632 | post_process_normal( |
| 633 | sampler_or_pooler_output, |
| 634 | model_output, |
| 635 | share_inputs, |
| 636 | sampling_metadata, |
| 637 | block_size, |
| 638 | think_end_id, |
| 639 | splitwise_role_is_decode, |
| 640 | enable_entropy, |
| 641 | routing_replay_manager, |
| 642 | ) |
| 643 | share_inputs["last_preempted_idx"].copy_(share_inputs["preempted_idx"]) |
| 644 | share_inputs["preempted_idx"][:] = 0 |
no test coverage detected