biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Fix attention

Browse Source
This commit is contained in:
Woosuk Kwon
2023-02-23 23:02:25 +00:00
parent ba84b8728a
commit 932844f1cd
2 changed files with 21 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
cacheflow/models/input_metadata.py
View File

@@ -11,6 +11,7 @@ class InputMetadata:
prompt_lens: List[int],
slot_mapping: torch.Tensor,
context_lens: torch.Tensor,
# FIXME: Rename
max_context_len: int,
block_tables: torch.Tensor,
) -> None:
@@ -23,9 +24,19 @@ class InputMetadata:
self.num_prompts = len(prompt_lens)
self.num_generation_tokens = context_lens.shape[0]
self.num_valid_tokens = len(slot_mapping)
if block_tables.numel() > 0:
self.max_num_blocks_per_seq = block_tables.shape[1]
else:
self.max_num_blocks_per_seq = 0
assert self.num_generation_tokens == block_tables.shape[0]
assert self.num_prompts + self.num_generation_tokens == len(seq_ids)
def __repr__(self) -> str:
return (f'InputMetadata('
f'seq_ids={self.seq_ids}, '
f'num_prompts={self.num_prompts}, '
f'num_generation_tokens={self.num_generation_tokens}, '
f'num_valid_tokens={self.num_valid_tokens}, '
f'max_num_blocks_per_seq={self.max_num_blocks_per_seq}, '
f'max_context_len={self.max_context_len})')