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

[BugFix] Reordering extend logic fix (#27739)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
This commit is contained in:
Lucas Wilkinson
2025-10-30 12:39:34 +08:00
committed by GitHub
parent b8c48c5d72
commit b5d70751d8
2 changed files with 23 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tests/v1/attention/test_batch_reordering.py
View File

@ -53,7 +53,7 @@ REORDER_TEST_CASES = {
expected_modified=True,
),
"already_ordered": ReorderTestCase(
requests=[(1, 10), (1, 20), (100, 100), (200, 200)],
requests=[(1, 10), (1, 20), (100, 100), (200, 0)],
expected_order=[0, 1, 2, 3],
expected_modified=False,
),
@ -74,15 +74,30 @@ REORDER_TEST_CASES = {
expected_modified=True,
),
"decode_extend_prefill": ReorderTestCase(
requests=[(100, 100), (10, 50), (1, 10)],
requests=[(100, 0), (10, 50), (1, 10)],
expected_order=[2, 1, 0],
expected_modified=True,
),
"extend_prefill_only": ReorderTestCase(
requests=[(100, 100), (10, 50), (200, 200), (20, 75)],
requests=[(100, 0), (10, 50), (200, 0), (20, 75)],
expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place
expected_modified=True,
),
"complicated_mixed_interleaved": ReorderTestCase(
requests=[
(1, 20),
(1, 50),
(374, 0),
(300, 20),
(1, 20),
(256, 0),
(1, 5),
(27, 0),
(1, 4),
],
expected_order=[0, 1, 6, 8, 4, 3, 2, 7, 5],
expected_modified=True,
),
}

10
vllm/v1/attention/backends/utils.py
View File

@ -811,8 +811,8 @@ def reorder_batch_to_split_decodes_and_prefills(
num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]
is_decode = num_scheduled_tokens_np <= decode_threshold
is_extend = (~is_decode) & (num_computed_tokens_np > num_scheduled_tokens_np)
is_prefill = (~is_decode) & (num_computed_tokens_np == num_scheduled_tokens_np)
is_extend = (~is_decode) & (num_computed_tokens_np > 0)
is_prefill = (~is_decode) & (num_computed_tokens_np == 0)
# Desired order: decode → extend → prefill
req_regions = np.zeros(is_decode.shape, dtype=np.int32) # 0 = decode by default
@ -832,11 +832,11 @@ def reorder_batch_to_split_decodes_and_prefills(
return False
# Extract indices that need swapping and sort by target region
swap_indices = np.where(needs_swap)[0]
orig_indices = np.where(needs_swap)[0]
sorted_order = np.argsort(req_regions[needs_swap], kind="stable")
dest_indices = swap_indices[sorted_order]
src_indices = orig_indices[sorted_order]
src_dest_map = {int(src): int(dst) for src, dst in zip(swap_indices, dest_indices)}
src_dest_map = {int(src): int(dst) for src, dst in zip(src_indices, orig_indices)}
for src in src_dest_map:
dst = src_dest_map[src]