Metadata for input sequences. Used in PagedAttention. Args: prompt_lens: Lengths of prompts. slot_mapping: The address to write the new KV to of each token. max_context_len: The maximum context length. context_lens: the length of attention context for each sequen
| 6 | |
| 7 | @dataclass |
| 8 | class InputMetadata: |
| 9 | """Metadata for input sequences. Used in PagedAttention. |
| 10 | |
| 11 | Args: |
| 12 | prompt_lens: Lengths of prompts. |
| 13 | slot_mapping: The address to write the new KV to of each token. |
| 14 | max_context_len: The maximum context length. |
| 15 | context_lens: the length of attention context for each sequence. |
| 16 | block_tables: The block tables. (Seq id -> list of physical block) |
| 17 | """ |
| 18 | |
| 19 | is_prompt: bool |
| 20 | slot_mapping: torch.Tensor |
| 21 | max_context_len: Optional[int] |
| 22 | context_lens: Optional[torch.Tensor] |
| 23 | block_tables: Optional[torch.Tensor] |
| 24 | max_seq_len: Optional[int] |
| 25 | seq_start_loc: Optional[torch.Tensor] |
| 26 | attn_bias = None |
| 27 | prompt_lens: Optional[List[int]] = None |
| 28 | decode_seq_lens: Optional[List[int]] = None |
| 29 | |
| 30 | use_flashinfer: bool = False |
| 31 | decode_wrapper = None |
| 32 | workspace_buffer: Optional[torch.Tensor] = None |
| 33 | # The indptr of the paged kv cache, shape: [batch_size + 1] |
| 34 | paged_kv_indptr: Optional[torch.Tensor] = None |
| 35 | # The page indices of the paged kv cache |
| 36 | paged_kv_indices: Optional[torch.Tensor] = None |
| 37 | # The number of entries in the last page of each request in |
| 38 | # the paged kv cache, shape: [batch_size] |
| 39 | paged_kv_last_page_len: Optional[torch.Tensor] = None |
| 40 | # The number of query/output heads |
| 41 | num_qo_heads: Optional[int] = None |
| 42 | # The number of key/value heads |
| 43 | num_kv_heads: Optional[int] = None |
| 44 | # The dimension of the attention heads |
| 45 | head_dim: Optional[int] = None |
| 46 | # Block size of vllm |
| 47 | page_size: Optional[int] = None |
| 48 | |
| 49 | def __post_init__(self): |
| 50 | if self.use_flashinfer and not self.is_prompt: |
| 51 | from flashinfer import BatchDecodeWithPagedKVCacheWrapper |
| 52 | self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( |
| 53 | self.workspace_buffer, "NHD") |
| 54 | self.decode_wrapper.end_forward() |
| 55 | self.decode_wrapper.begin_forward( |
| 56 | self.paged_kv_indptr, |
| 57 | self.paged_kv_indices, |
| 58 | self.paged_kv_last_page_len, |
| 59 | self.num_qo_heads, |
| 60 | self.num_kv_heads, |
| 61 | self.head_dim, |
| 62 | self.page_size, |
| 63 | # Disable flashinfer's pos encoding and use vllm's rope. |
| 64 | pos_encoding_mode="NONE") |
| 65 |
no outgoing calls
no test coverage detected