A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 01:45 UTC
[BugFix] Reordering extend logic fix (#27739)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Oct 30, 2025 at 12:39 UTC
b5d70751d82c272a72f105299ef24ae316c41ded
2 files changed
+23
-8
tests/v1/attention/test_batch_reordering.py
+18
-3
@@ -53,7 +53,7 @@ REORDER_TEST_CASES = {
53
expected_modified=True,
54
),
55
"already_ordered": ReorderTestCase(
56
- requests=[(1, 10), (1, 20), (100, 100), (200, 200)],
56
+ requests=[(1, 10), (1, 20), (100, 100), (200, 0)],
57
expected_order=[0, 1, 2, 3],
58
expected_modified=False,
59
),
@@ -74,15 +74,30 @@ REORDER_TEST_CASES = {
74
expected_modified=True,
75
),
76
"decode_extend_prefill": ReorderTestCase(
77
- requests=[(100, 100), (10, 50), (1, 10)],
77
+ requests=[(100, 0), (10, 50), (1, 10)],
78
expected_order=[2, 1, 0],
79
expected_modified=True,
80
),
81
"extend_prefill_only": ReorderTestCase(
82
- requests=[(100, 100), (10, 50), (200, 200), (20, 75)],
82
+ requests=[(100, 0), (10, 50), (200, 0), (20, 75)],
83
expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place
84
expected_modified=True,
85
),
86
+ "complicated_mixed_interleaved": ReorderTestCase(
87
+ requests=[
88
+ (1, 20),
89
+ (1, 50),
90
+ (374, 0),
91
+ (300, 20),
92
+ (1, 20),
93
+ (256, 0),
94
+ (1, 5),
95
+ (27, 0),
96
+ (1, 4),
97
+ ],
98
+ expected_order=[0, 1, 6, 8, 4, 3, 2, 7, 5],
99
+ expected_modified=True,
100
+ ),
101
}
102
103
vllm/v1/attention/backends/utils.py
+5
-5
@@ -811,8 +811,8 @@ def reorder_batch_to_split_decodes_and_prefills(
811
num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]
812
813
is_decode = num_scheduled_tokens_np <= decode_threshold
814
- is_extend = (~is_decode) & (num_computed_tokens_np > num_scheduled_tokens_np)
815
- is_prefill = (~is_decode) & (num_computed_tokens_np == num_scheduled_tokens_np)
814
+ is_extend = (~is_decode) & (num_computed_tokens_np > 0)
815
+ is_prefill = (~is_decode) & (num_computed_tokens_np == 0)
816
817
# Desired order: decode → extend → prefill
818
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(
832
return False
833
834
# Extract indices that need swapping and sort by target region
835
- swap_indices = np.where(needs_swap)[0]
835
+ orig_indices = np.where(needs_swap)[0]
836
sorted_order = np.argsort(req_regions[needs_swap], kind="stable")
837
- dest_indices = swap_indices[sorted_order]
837
+ src_indices = orig_indices[sorted_order]
838
839
- src_dest_map = {int(src): int(dst) for src, dst in zip(swap_indices, dest_indices)}
839
+ src_dest_map = {int(src): int(dst) for src, dst in zip(src_indices, orig_indices)}
840
841
for src in src_dest_map:
842
dst = src_dest_map[src]