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, 05:30 UTC
[Attention] Remove unused reorder_batch method (#24463)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 6, 2025 at 13:13 UTC
4727a8afa79528480212aa537290f4dce638d1ee
6 files changed
+8
-64
tests/v1/logits_processors/test_correctness.py
+1
-1
@@ -581,7 +581,7 @@ def _generate_fake_step_update(
581
persistent_batch[:] = persistent_batch[0:condensed_batch_size]
582
583
if condensed_batch_size > 1:
584
- # Simulate arbitrary reorder_batch() in the kernel backend
584
+ # Simulate arbitrary batch ordering in the kernel backend
585
# Generate a random number k of non-overlapping swap tuples
586
k = random.randint(0, condensed_batch_size // 2)
587
idxs = list(range(condensed_batch_size))
vllm/v1/attention/backends/flashinfer.py
+2
-4
@@ -602,8 +602,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
602
)
603
else:
604
# Regular attention (common case).
605
- # Decodes are at the front and prefills are at the back,
606
- # according to reorder_batch()
605
+ # Decodes are at the front and prefills are at the back.
606
num_prefills = attn_metadata.num_prefills
607
num_decodes = attn_metadata.num_decodes
608
if num_prefills > 0:
@@ -925,8 +924,7 @@ class FlashInferImpl(AttentionImpl):
924
stride_order = FlashInferBackend.get_kv_cache_stride_order()
925
kv_cache_permute = kv_cache.permute(*stride_order)
926
# Regular attention (common case).
928
- # Decodes are at the front and prefills are at the back,
929
- # according to reorder_batch()
927
+ # Decodes are at the front and prefills are at the back.
928
if num_prefill_tokens > 0:
929
prefill_wrapper = attn_metadata.prefill_wrapper
930
prefill_query = query[num_decode_tokens:]
vllm/v1/attention/backends/flex_attention.py
+1
-10
@@ -3,7 +3,7 @@
3
"""Attention layer with FlexAttention."""
4
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Optional, Union
6
+from typing import Optional, Union
7
8
import torch
9
import torch._dynamo.decorators
@@ -38,10 +38,6 @@ from vllm.v1.kv_cache_interface import AttentionSpec
38
39
logger = init_logger(__name__)
40
41
-if TYPE_CHECKING:
42
- from vllm.v1.core.sched.output import SchedulerOutput
43
- from vllm.v1.worker.gpu_input_batch import InputBatch
44
-
41
create_block_mask_compiled = torch.compile(
42
create_block_mask, fullgraph=True, mode="reduce-overhead"
43
)
@@ -600,11 +596,6 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
596
self.q_block_size: int = 16 if is_torch_equal_or_newer("2.9.0.dev0") else 128
597
self.kv_block_size: int = 16 if is_torch_equal_or_newer("2.9.0.dev0") else 128
598
603
- def reorder_batch(
604
- self, input_batch: "InputBatch", scheduler_output: "SchedulerOutput"
605
- ) -> bool:
606
- return False
607
-
599
def build(
600
self,
601
common_prefix_len: int,
vllm/v1/attention/backends/tree_attn.py
+3
-14
@@ -4,10 +4,11 @@
4
5
import ast
6
from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, Optional
7
+from typing import Optional
8
9
import torch
10
11
+from vllm import _custom_ops as ops
12
from vllm.attention.backends.abstract import (
13
AttentionBackend,
14
AttentionImpl,
@@ -20,17 +21,10 @@ from vllm.logger import init_logger
21
from vllm.v1.attention.backends.utils import (
22
AttentionMetadataBuilder,
23
CommonAttentionMetadata,
23
- reorder_batch_to_split_decodes_and_prefills,
24
split_decodes_and_prefills,
25
)
26
from vllm.v1.kv_cache_interface import AttentionSpec
27
28
-if TYPE_CHECKING:
29
- from vllm.v1.core.sched.output import SchedulerOutput
30
- from vllm.v1.worker.gpu_input_batch import InputBatch
31
-
32
-from vllm import _custom_ops as ops
33
-
28
logger = init_logger(__name__)
29
30
@@ -189,12 +183,7 @@ class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadat
183
device=device,
184
)
185
192
- def reorder_batch(
193
- self, input_batch: "InputBatch", scheduler_output: "SchedulerOutput"
194
- ) -> bool:
195
- return reorder_batch_to_split_decodes_and_prefills(
196
- input_batch, scheduler_output, decode_threshold=self.tree_attn_bias.shape[0]
197
- )
186
+ self.reorder_batch_threshold = self.tree_attn_bias.shape[0]
187
188
def build(
189
self,
vllm/v1/attention/backends/utils.py
-22
@@ -299,24 +299,6 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
299
"""
300
raise NotImplementedError
301
302
- def reorder_batch(
303
- self, input_batch: "InputBatch", scheduler_output: "SchedulerOutput"
304
- ) -> bool:
305
- """
306
- Update the order of requests in the batch based on the attention
307
- backend's needs. For example, some attention backends (namely MLA) may
308
- want to separate requests based on if the attention computation will be
309
- compute-bound or memory-bound.
310
-
311
- Args:
312
- input_batch: input batch
313
- scheduler_output: scheduler output.
314
-
315
- Returns:
316
- True if the batch was modified, False otherwise.
317
- """
318
- raise NotImplementedError
319
-
302
def build_for_cudagraph_capture(
303
self, common_attn_metadata: CommonAttentionMetadata
304
) -> M:
@@ -828,10 +810,6 @@ def reorder_batch_to_split_decodes_and_prefills(
810
811
for i, req_id in enumerate(input_batch.req_ids):
812
num_tokens = scheduler_output.num_scheduled_tokens[req_id]
831
- # for now treat 1 scheduled token as "decode" even if it's not,
832
- # we should update this to something like < 8 in the future but
833
- # currently the TritonMLA._forward_decode only supports
834
- # num_tokens = 1
813
if num_tokens <= decode_threshold:
814
decodes.append(i)
815
num_decode_tokens += num_tokens
vllm/v1/attention/backends/xformers.py
+1
-13
@@ -3,7 +3,7 @@
3
"""Attention layer with XFormersAttention."""
4
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Optional
6
+from typing import Optional
7
8
import torch
9
@@ -19,7 +19,6 @@ from vllm.logger import init_logger
19
from vllm.v1.attention.backends.utils import (
20
AttentionMetadataBuilder,
21
CommonAttentionMetadata,
22
- reorder_batch_to_split_decodes_and_prefills,
22
split_decodes_and_prefills,
23
)
24
from vllm.v1.kv_cache_interface import AttentionSpec
@@ -35,10 +34,6 @@ try:
34
except ImportError:
35
XFORMERS_AVAILABLE = False
36
38
-if TYPE_CHECKING:
39
- from vllm.v1.core.sched.output import SchedulerOutput
40
- from vllm.v1.worker.gpu_input_batch import InputBatch
41
-
37
from vllm import _custom_ops as ops
38
39
logger = init_logger(__name__)
@@ -223,13 +218,6 @@ class XFormersAttentionMetadataBuilder(
218
self._num_decodes = 0
219
self._num_decode_tokens = 0
220
226
- def reorder_batch(
227
- self, input_batch: "InputBatch", scheduler_output: "SchedulerOutput"
228
- ) -> bool:
229
- return reorder_batch_to_split_decodes_and_prefills(
230
- input_batch, scheduler_output, decode_threshold=self.reorder_batch_threshold
231
- )
232
-
221
def build(
222
self,
223
common_prefix_len: int,