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 11, 18:45 UTC
[Attention] Make `split_decodes_and_prefills(..., require_uniform=True)` support padding (#29644)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Lucas Wilkinson committed
Dec 9, 2025 at 02:24 UTC
aed846917fee3b27df876d21ade4d3a30d4de402
2 files changed
+31
-4
tests/v1/attention/test_attention_splitting.py
+24
-1
@@ -154,7 +154,10 @@ def test_split_attn_metadata_decode_batch(large_decode_metadata):
154
155
156
def apply_split_decodes_and_prefills(
157
- query_lens: list[int], decode_threshold: int, require_uniform: bool
157
+ query_lens: list[int],
158
+ decode_threshold: int,
159
+ require_uniform: bool,
160
+ padded_num_tokens: int | None = None,
161
):
162
"""Helper function to apply split_decodes_and_prefills and return
163
the results."""
@@ -165,6 +168,10 @@ def apply_split_decodes_and_prefills(
168
block_size=16,
169
device=device,
170
)
171
+
172
+ if padded_num_tokens is not None:
173
+ common_metadata.num_actual_tokens = padded_num_tokens
174
+
175
return split_decodes_and_prefills(
176
common_metadata,
177
decode_threshold=decode_threshold,
@@ -271,6 +278,22 @@ def test_split_decodes_and_prefills_uniform_mixed_batch_non_uniform_decodes():
278
assert num_prefill_tokens == (sum(query_lens) - 2) # rest of the tokens
279
280
281
+def test_split_decodes_and_prefills_uniform_padded_batch_all_same():
282
+ """uniform batch where all query lengths are identical with 0 length padded reqs."""
283
+ # All query lengths are 2, with decode_threshold=3 (so 2 <= 3)
284
+ # This triggers the padded uniform path at line 891
285
+ query_lens = [2, 2, 2, 0]
286
+ padded_num_tokens = 8
287
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
288
+ apply_split_decodes_and_prefills(query_lens, 3, True, padded_num_tokens)
289
+ )
290
+ # With uniform batch, all requests are treated as decodes
291
+ assert num_decodes == 4
292
+ assert num_prefills == 0
293
+ assert num_decode_tokens == padded_num_tokens
294
+ assert num_prefill_tokens == 0
295
+
296
+
297
@pytest.mark.parametrize(
298
"seq_lens,query_lens,split_point,expected_first_reqs,expected_second_reqs",
299
[
vllm/v1/attention/backends/utils.py
+7
-3
@@ -883,11 +883,15 @@ def split_decodes_and_prefills(
883
return 0, num_reqs, 0, num_tokens
884
885
if require_uniform:
886
+ # check if we are in a padded uniform batch; this is used for full-CGs, some
887
+ # requests may have a query length of 0 but since they are padding its fine
888
+ # to treat them as decodes (ensures num_decodes matches the captured size)
889
+ if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
890
+ assert num_reqs * query_lens[0] == num_tokens, "tokens not padded correctly"
891
+ return num_reqs, 0, num_tokens, 0 # all decodes
892
is_prefill = query_lens != query_lens[0]
893
else:
888
- # 0-query len indicates a padded request; leave this at the back
889
- # of the batch with the prefills
890
- is_prefill = (query_lens > decode_threshold) | (query_lens == 0)
894
+ is_prefill = query_lens > decode_threshold
895
896
if not torch.any(is_prefill):
897
return num_reqs, 0, num_tokens, 0