| 1792 | return num_blocks, max_blocks_per_seq |
| 1793 | |
| 1794 | def setup(self, |
| 1795 | batch_size: int, |
| 1796 | max_context_length: int, |
| 1797 | max_new_tokens: int, |
| 1798 | beam_width: int = 1, |
| 1799 | max_attention_window_size: Optional[int] = None, |
| 1800 | sink_token_length: Optional[int] = None, |
| 1801 | encoder_max_input_length: Optional[int] = None, |
| 1802 | lora_manager: LoraManager = None, |
| 1803 | lora_uids: List[str] = None, |
| 1804 | medusa_choices: List[List[int]] = None, |
| 1805 | multi_block_mode: bool = True, |
| 1806 | enable_context_fmha_fp32_acc: bool = None): |
| 1807 | # Store these params related to buffer size to check against |
| 1808 | # the input shape with the params given in decode() |
| 1809 | self.batch_size = batch_size |
| 1810 | self.max_context_length = max_context_length |
| 1811 | self.max_new_tokens = max_new_tokens |
| 1812 | self.max_seq_length = max_context_length + max_new_tokens |
| 1813 | if medusa_choices is not None or self.is_redrafter_mode: |
| 1814 | self.max_seq_length += self.max_draft_tokens |
| 1815 | self.beam_width = beam_width |
| 1816 | self.encoder_max_input_length = encoder_max_input_length |
| 1817 | self.multi_block_mode = multi_block_mode |
| 1818 | self.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc |
| 1819 | if max_attention_window_size is None: |
| 1820 | self.max_attention_window_size = self.max_seq_length |
| 1821 | logger.debug( |
| 1822 | "The max_attention_window_size is not set, we will use max_seq_length by default." |
| 1823 | ) |
| 1824 | self.host_max_attention_window_sizes = torch.ones( |
| 1825 | (self.num_attn_layers, ), |
| 1826 | dtype=torch.int32) * self.max_attention_window_size |
| 1827 | |
| 1828 | elif isinstance(max_attention_window_size, int): |
| 1829 | if max_attention_window_size > self.max_seq_length: |
| 1830 | logger.warning( |
| 1831 | "The value of max_attention_window_size should ideally not exceed max_seq_length. " |
| 1832 | "Therefore, it has been adjusted to match the value of max_seq_length." |
| 1833 | ) |
| 1834 | self.max_attention_window_size = min(max_attention_window_size, |
| 1835 | self.max_seq_length) |
| 1836 | self.host_max_attention_window_sizes = torch.ones( |
| 1837 | (self.num_attn_layers, ), |
| 1838 | dtype=torch.int32) * self.max_attention_window_size |
| 1839 | |
| 1840 | elif isinstance(max_attention_window_size, (torch.Tensor, list)): |
| 1841 | if isinstance(max_attention_window_size, list): |
| 1842 | max_attention_window_size = torch.tensor( |
| 1843 | max_attention_window_size, dtype=torch.int32) |
| 1844 | self.max_attention_window_size = int( |
| 1845 | torch.max(max_attention_window_size).item()) |
| 1846 | attn_win_size_len = max_attention_window_size.shape[0] |
| 1847 | num_total_attn_layers = self.layer_types.count('attention') |
| 1848 | if attn_win_size_len < num_total_attn_layers: |
| 1849 | repeat_num = num_total_attn_layers // attn_win_size_len |
| 1850 | remain_num = num_total_attn_layers % attn_win_size_len |
| 1851 | warning_info = "The size of max_attention_window_size tensor/list is less than num_attn_layers, " \ |