Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix blocks number calculation for Flat PA #269

Merged
merged 3 commits into from
Sep 18, 2024
Merged

Fix blocks number calculation for Flat PA #269

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a4e3722
Fix blocks number calculation for Flat PA
iboiko-habana Sep 11, 2024
041659d
Add format check
iboiko-habana Sep 11, 2024
0907f38
Update habana_model_runner.py
iboiko-habana Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 16 additions & 8 deletions vllm/worker/habana_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,11 +171,16 @@ def generate_prompt_buckets(bs_bucket_config,
def generate_decode_buckets(bs_bucket_config, blocks_bucket_config,
max_blocks):
buckets = []
for bs in warmup_range(bs_bucket_config):
for blocks in warmup_range(blocks_bucket_config):
bs_buckets = warmup_range(bs_bucket_config)
block_buckets = warmup_range(blocks_bucket_config)
bmin, bstep, bmax = blocks_bucket_config
last_bucket = max_blocks if (max_blocks // bstep
== 0) else (max_blocks // bstep + 1) * bstep
for bs in bs_buckets:
for blocks in block_buckets:
if blocks < bs:
continue
if blocks > max_blocks:
if blocks > last_bucket:
break
buckets.append((bs, blocks))
return list(sorted(buckets, key=lambda b: (b[0] * b[1], b[1], b[0])))
Expand Down Expand Up @@ -918,10 +923,12 @@ def _prepare_decode(
seq_lens.append(seq_len)

block_table = seq_group_metadata.block_tables[seq_id]
block_number = block_table[position // self.block_size]
if block_number == _PAD_BLOCK_ID:
if len(block_table) == 0:
block_number = _PAD_BLOCK_ID
block_table = []
slot = next(dummy_slots)
else:
block_number = block_table[position // self.block_size]
block_offset = position % self.block_size
slot = block_number * self.block_size + block_offset
slot_mapping.append([slot])
Expand All @@ -946,7 +953,7 @@ def _prepare_decode(

num_decode_tokens = sum(seq_lens)

blocks_used = [len(bt) for bt in block_tables]
blocks_used = [len(bt) for bt in block_tables if bt]
block_list = list(itertools.chain(*block_tables))
block_mapping_nested: List[List[int]] = [
[i] * b_u for i, b_u in enumerate(blocks_used)
Expand Down Expand Up @@ -1034,8 +1041,9 @@ def prepare_input_tensors(
batch_size_padded = find_bucket(real_batch_size, bucket_cfg)
batch_size_padding = batch_size_padded - real_batch_size
seq_group_metadata_list = seq_group_metadata_list.copy()
seq_group_metadata_list.extend(seq_group_metadata_list[0]
for _ in range(batch_size_padding))
seq_group_metadata_list.extend(
self.create_dummy_seq_group_metadata(0, 0, is_prompt)
for _ in range(batch_size_padding))

prefill_reqs = []
decode_reqs = []
Expand Down
Loading