Skip to content

GitLab

Commit 7f985166 authored by Woosuk Kwon's avatar Woosuk Kwon
Browse files

Consider pempty tensor

parent 86f9eb6d
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 1 deletion
cacheflow/models/input_metadata.py
View file @ 7f985166
...@@ -23,6 +23,9 @@ class InputMetadata: ...@@ -23,6 +23,9 @@ class InputMetadata:
self.num_prompts = len(prompt_lens) self.num_prompts = len(prompt_lens)
self.num_generation_tokens = context_lens.shape[0] self.num_generation_tokens = context_lens.shape[0]
if block_tables.numel() > 0:
self.max_num_blocks_per_seq = block_tables.shape[1] 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_generation_tokens == block_tables.shape[0]
assert self.num_prompts + self.num_generation_tokens == len(seq_ids) assert self.num_prompts + self.num_generation_tokens == len(seq_ids)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment