TODO(gongshaotian): normalization name
(
share_inputs: InputBatch,
block_size: int,
enc_dec_block_num: int,
speculative_config: SpeculativeConfig,
enable_prefix_caching: bool = False,
)
| 645 | |
| 646 | |
| 647 | def step_cuda( |
| 648 | share_inputs: InputBatch, |
| 649 | block_size: int, |
| 650 | enc_dec_block_num: int, |
| 651 | speculative_config: SpeculativeConfig, |
| 652 | enable_prefix_caching: bool = False, |
| 653 | ) -> None: |
| 654 | """ |
| 655 | TODO(gongshaotian): normalization name |
| 656 | """ |
| 657 | |
| 658 | if speculative_config.method is not None: |
| 659 | if DISABLE_RECOVER: |
| 660 | speculate_step_reschedule( |
| 661 | share_inputs["stop_flags"], |
| 662 | share_inputs["seq_lens_this_time"], |
| 663 | share_inputs["step_seq_lens_encoder"], |
| 664 | share_inputs["seq_lens_encoder"], |
| 665 | share_inputs["seq_lens_decoder"], |
| 666 | share_inputs["block_tables"], |
| 667 | share_inputs["encoder_block_lens"], |
| 668 | share_inputs["is_block_step"], |
| 669 | share_inputs["step_block_list"], |
| 670 | share_inputs["step_lens"], |
| 671 | share_inputs["recover_block_list"], |
| 672 | share_inputs["recover_lens"], |
| 673 | share_inputs["need_block_list"], |
| 674 | share_inputs["need_block_len"], |
| 675 | share_inputs["used_list_len"], |
| 676 | share_inputs["free_list"], |
| 677 | share_inputs["free_list_len"], |
| 678 | share_inputs["input_ids"], |
| 679 | share_inputs["pre_ids"], |
| 680 | share_inputs["step_idx"], |
| 681 | share_inputs["next_tokens"], |
| 682 | share_inputs["first_token_ids"], |
| 683 | share_inputs["accept_num"], |
| 684 | block_size, |
| 685 | enc_dec_block_num, |
| 686 | speculative_config.num_speculative_tokens, |
| 687 | ) |
| 688 | else: |
| 689 | if enable_prefix_caching: |
| 690 | speculate_step_system_cache( |
| 691 | share_inputs["stop_flags"], |
| 692 | share_inputs["seq_lens_this_time"], |
| 693 | share_inputs["step_seq_lens_encoder"], |
| 694 | share_inputs["step_seq_lens_decoder"], |
| 695 | share_inputs["seq_lens_encoder"], |
| 696 | share_inputs["seq_lens_decoder"], |
| 697 | share_inputs["block_tables"], |
| 698 | share_inputs["encoder_block_lens"], |
| 699 | share_inputs["is_block_step"], |
| 700 | share_inputs["step_block_list"], |
| 701 | share_inputs["step_lens"], |
| 702 | share_inputs["recover_block_list"], |
| 703 | share_inputs["recover_lens"], |
| 704 | share_inputs["need_block_list"], |
no outgoing calls
no test coverage detected