Skip to content

[0.10.2] [SW-241908] Omit all prompt buckets that exceed max_num_batched_tokens (#331) - cherry-pick #342

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

Open
wants to merge 5 commits into
base: v0.10.2_next
Choose a base branch
from
Open

[0.10.2] [SW-241908] Omit all prompt buckets that exceed max_num_batched_tokens (#331) - cherry-pick #342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
23e02ff
Omit all prompt buckets that exceed max_num_batched_tokens
skavulya Oct 7, 2025
b3c56e4
Update prompt bucket filter criteria
skavulya Oct 7, 2025
5f40882
Merge branch 'v0.10.2_next' into skavulya/prompt_bucket_fix_v0.10.2_next
michalkuligowski Oct 9, 2025
cb6db1a
Merge branch 'v0.10.2_next' into skavulya/prompt_bucket_fix_v0.10.2_next
skavulya Oct 13, 2025
1b98f0a
Update prompt bucketing test
skavulya Oct 14, 2025
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
2 changes: 1 addition & 1 deletion tests/unit_tests/test_bucketing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def test_generate_prompt_buckets():
ctx_range = [0, 1, 2, 3, 4]
buckets = generate_buckets(bs_range, query_range, ctx_range, True, max_model_len, bs, prompt_bs,
max_num_batched_tokens, block_size, max_blocks)
assert len(buckets) == 40
assert len(buckets) == 17


def test_generate_decode_buckets():
Expand Down
21 changes: 11 additions & 10 deletions vllm_gaudi/extension/bucketing/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,22 +243,23 @@ def expand_to_neighbor_buckets(bs_idx, bs_range, query_idx, query_range, max_num
'''
candidates = [(bs_idx, query_idx), (bs_idx + 1, query_idx), (bs_idx, query_idx + 1),
(bs_idx + 1, query_idx + 1)]
valid = bs_range[bs_idx] * query_range[query_idx] <= max_num_batched_tokens
if not valid:
omitted_buckets.add(("bs_range[bs_idx] * query_range[query_idx] <= max_num_batched_tokens",
"-> bs, quesry: ", bs_idx, query_idx))
return {}
valid_candidates = [(b_idx, q_idx) for b_idx, q_idx in candidates
if b_idx < len(bs_range) and q_idx < len(query_range)]
if (b_idx < len(bs_range) and q_idx < len(query_range))]
return {(bs_range[b_idx], query_range[q_idx]) for b_idx, q_idx in valid_candidates}

# filter rules for buckets
# prompt
def not_over_max_model_len(bs, query, ctx):
if not query + ctx * block_size <= max_model_len:
if not bs * (query + ctx * block_size) <= max_model_len:
omitted_buckets.add(
("condition: query + ctx * block_size <= max_model_len", "-> bs, quesry, ctx: ", bs, query, ctx))
return query + ctx * block_size <= max_model_len
("condition: bs * (query + ctx * block_size) <= max_model_len", "-> bs, query, ctx: ", bs, query, ctx))
return bs * (query + ctx * block_size) <= max_model_len

def not_over_max_num_batched_tokens(bs, query, ctx):
if not bs * query <= max_num_batched_tokens:
omitted_buckets.add(
("condition: bs * query <= max_num_batched_tokens", "-> bs, query, ctx: ", bs, query, ctx))
return bs * query <= max_num_batched_tokens

def ctx_not_over_max_ctx_for_merged_prefill(bs, query, ctx):
if not ctx <= max_num_prefill_seqs * math.ceil(
Expand All @@ -285,7 +286,7 @@ def batch_size_smaller_than_blocks(bs, query, ctx):
"prompt": {
# depends only on merged_prefill
True: [ctx_not_over_max_ctx_for_merged_prefill],
False: [not_over_max_model_len],
False: [not_over_max_model_len, not_over_max_num_batched_tokens],
},
"decode": {
# depends only on contiguous PA
Expand Down