@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4abfd8796f

Fork 0

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, 22:35 UTC

[V1] [Hybrid] Validate compatibility of attention backend batch reordering at init time (#21557)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>

Thomas Parnell committed Aug 2, 2025 at 14:29 UTC 4abfd8796f37adc8fccc9481f37f20de1bce62e4

7 files changed +96 -72

vllm/v1/attention/backends/flashinfer.py

+12 -16

@@ -4,7 +4,7 @@

4 from __future__ import annotations

6 from dataclasses import dataclass

7 -from typing import TYPE_CHECKING, ClassVar, Optional, Union

7 +from typing import ClassVar, Optional, Union

9 import torch

10 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,

@@ -21,17 +21,17 @@ from vllm.logger import init_logger

21 from vllm.utils import cdiv, is_pin_memory_available

22 from vllm.utils.flashinfer import use_trtllm_decode_attention

23 from vllm.v1.attention.backends.flash_attn import use_cascade_attention

24 -from vllm.v1.attention.backends.utils import (

25 - AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,

26 - get_kv_cache_layout, get_per_layer_parameters,

27 - infer_global_hyperparameters, reorder_batch_to_split_decodes_and_prefills,

28 - split_decodes_and_prefills)

24 +# yapf conflicts with isort for this block

25 +# yapf: disable

26 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,

27 + AttentionMetadataBuilder,

28 + CommonAttentionMetadata,

29 + get_kv_cache_layout,

30 + get_per_layer_parameters,

31 + infer_global_hyperparameters,

32 + split_decodes_and_prefills)

33 from vllm.v1.kv_cache_interface import AttentionSpec

31 -if TYPE_CHECKING:

32 - from vllm.v1.core.sched.output import SchedulerOutput

33 - from vllm.v1.worker.gpu_input_batch import InputBatch

34 -

35 FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024

37 logger = init_logger(__name__)

@@ -179,6 +179,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

179 attn_cudagraph_support: ClassVar[AttentionCGSupport] = \

180 AttentionCGSupport.PURE_DECODE_ONLY

181

182 + reorder_batch_threshold: ClassVar[int] = 1

183 +

184 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

185 vllm_config: VllmConfig, device: torch.device):

186 self.device = device

@@ -239,12 +241,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

241 dtype=torch.int32,

242 device=self.device)

243

242 - def reorder_batch(self, input_batch: InputBatch,

243 - scheduler_output: SchedulerOutput) -> bool:

244 - return reorder_batch_to_split_decodes_and_prefills(input_batch,

245 - scheduler_output,

246 - decode_threshold=1)

247 -

244 def _get_workspace_buffer(self):

245 if self._workspace_buffer is None:

246 self._workspace_buffer = torch.empty(

vllm/v1/attention/backends/mamba_attn.py

+6 -14

@@ -2,21 +2,17 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from dataclasses import dataclass

5 -from typing import TYPE_CHECKING, Optional

5 +from typing import ClassVar, Optional

7 import torch

9 from vllm.attention.backends.abstract import AttentionBackend

10 from vllm.config import VllmConfig

11 -from vllm.v1.attention.backends.utils import (

12 - AttentionMetadataBuilder, CommonAttentionMetadata,

13 - reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)

11 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

12 + CommonAttentionMetadata,

13 + split_decodes_and_prefills)

14 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec

16 -if TYPE_CHECKING:

17 - from vllm.v1.core.sched.output import SchedulerOutput

18 - from vllm.v1.worker.gpu_input_batch import InputBatch

19 -

17 def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,

18 chunk_size: int,

@@ -87,6 +83,8 @@ class Mamba2AttentionMetadata:

83 class Mamba2AttentionMetadataBuilder(

84 AttentionMetadataBuilder[Mamba2AttentionMetadata]):

86 + reorder_batch_threshold: ClassVar[int] = 1

87 +

88 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

89 vllm_config: VllmConfig, device: torch.device):

90 assert isinstance(kv_cache_spec, MambaSpec)

@@ -95,12 +93,6 @@ class Mamba2AttentionMetadataBuilder(

93 assert self.chunk_size is not None, (

94 "chunk_size needs to be set in the model config for Mamba2 models")

98 - def reorder_batch(self, input_batch: "InputBatch",

99 - scheduler_output: "SchedulerOutput") -> bool:

100 - return reorder_batch_to_split_decodes_and_prefills(input_batch,

101 - scheduler_output,

102 - decode_threshold=1)

103 -

96 def build(self,

97 common_prefix_len: int,

98 common_attn_metadata: CommonAttentionMetadata,

vllm/v1/attention/backends/mla/common.py

+7 -15

@@ -190,7 +190,7 @@ return curr_o @ W_O

190 import functools

191 from abc import abstractmethod

192 from dataclasses import dataclass, field

193 -from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union

193 +from typing import ClassVar, Generic, Optional, TypeVar, Union

194

195 import torch

196

@@ -210,10 +210,11 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

210 from vllm.platforms import current_platform

211 from vllm.utils import cdiv, round_down

212 from vllm.utils.flashinfer import has_nvidia_artifactory

213 -from vllm.v1.attention.backends.utils import (

214 - AttentionMetadataBuilder, CommonAttentionMetadata,

215 - get_per_layer_parameters, infer_global_hyperparameters,

216 - reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)

213 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

214 + CommonAttentionMetadata,

215 + get_per_layer_parameters,

216 + infer_global_hyperparameters,

217 + split_decodes_and_prefills)

218 from vllm.v1.kv_cache_interface import AttentionSpec

219

220 try:

@@ -233,10 +234,6 @@ try:

234 except ImportError:

235 flashinfer_available = False

236

236 -if TYPE_CHECKING:

237 - from vllm.v1.core.sched.output import SchedulerOutput

238 - from vllm.v1.worker.gpu_input_batch import InputBatch

239 -

237 logger = init_logger(__name__)

238

239 CUDNN_WORKSPACE_SIZE = 12800

@@ -403,6 +400,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

400 NOTE: Please read the comment at the top of the file before trying to

401 understand this class

402 """

403 + reorder_batch_threshold: ClassVar[int] = 1

404

405 def __init__(self,

406 kv_cache_spec: AttentionSpec,

@@ -559,12 +557,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

557 prefill.prefill_main = self._fi_prefill_main

558 prefill.prefill_chunks = self._fi_prefill_chunks

559

562 - def reorder_batch(self, input_batch: "InputBatch",

563 - scheduler_output: "SchedulerOutput") -> bool:

564 - return reorder_batch_to_split_decodes_and_prefills(input_batch,

565 - scheduler_output,

566 - decode_threshold=1)

567 -

560 def _build_decode(self, block_table_tensor: torch.Tensor,

561 seq_lens: torch.Tensor):

562 return MLACommonDecodeMetadata(

vllm/v1/attention/backends/rocm_aiter_fa.py

-3

@@ -251,9 +251,6 @@ class AiterFlashAttentionMetadataBuilder(

251 self.aot_sliding_window: Optional[tuple[int, int]] = None

252 self.total_tokens: int = 0

253

254 - def reorder_batch(self, input_batch, scheduler_output) -> bool:

255 - return False

256 -

254 def build_for_cudagraph_capture(

255 self, common_attn_metadata: CommonAttentionMetadata):

256 self.total_tokens = self.model_config.max_model_len \

vllm/v1/attention/backends/utils.py

+4 -8

@@ -167,6 +167,10 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

167 # Does this backend/builder support CUDA Graphs for attention.

168 attn_cudagraph_support: ClassVar[AttentionCGSupport] = \

169 AttentionCGSupport.NEVER

170 + # Does this backend/builder reorder the batch?

171 + # If not, set this to None. Otherwise set it to the query

172 + # length that will be pulled into the front of the batch.

173 + reorder_batch_threshold: ClassVar[Optional[int]] = None

174

175 @abstractmethod

176 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

@@ -221,14 +225,6 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

225 ) -> bool:

226 return False

227

224 - def reorder_batch(self, input_batch: "InputBatch",

225 - scheduler_output: "SchedulerOutput") -> bool:

226 - """

227 - This method can reorder the batch if desired by the backend.

228 - :return: Has the batch been reordered (default False).

229 - """

230 - return False

231 -

228

229 @functools.lru_cache

230 def get_kv_cache_layout():

vllm/v1/worker/cpu_model_runner.py

+33 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from contextlib import contextmanager

4 -from typing import Any

4 +from typing import TYPE_CHECKING, Any

6 import torch

7 import torch.nn as nn

@@ -9,8 +9,12 @@ import torch.nn as nn

9 from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 from vllm.model_executor.model_loader import get_model

12 +from vllm.v1.attention.backends.cpu_attn import TorchSDPAMetadataBuilderV1

13 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

15 +if TYPE_CHECKING:

16 + from vllm.v1.core.sched.output import SchedulerOutput

17 +

18 logger = init_logger(__name__)

@@ -27,6 +31,34 @@ class CPUModelRunner(GPUModelRunner):

32 self._postprocess_tenosrs()

34 + def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:

35 + """

36 + Update the order of requests in the batch based on the attention

37 + backend's needs. For example, some attention backends (namely MLA) may

38 + want to separate requests based on if the attention computation will be

39 + compute-bound or memory-bound.

40 +

41 + Args:

42 + scheduler_output: The scheduler output.

43 + """

44 + # Attention free models have zero kv_cache_goups, however models

45 + # like Mamba are also attention free but use the kv_cache for

46 + # keeping its internal state. This is why we check the number

47 + # of kv_cache groups instead of solely checking

48 + # for self.model_config.is_attention_free.

49 + if len(self.kv_cache_config.kv_cache_groups) == 0:

50 + return

51 +

52 + if len(self.kv_cache_config.kv_cache_groups) > 1:

53 + raise ValueError("Multiple KVCacheGroups is not"

54 + "currently supported with CPU model runner.")

55 +

56 + assert type(

57 + self.attn_metadata_builders[0]) is TorchSDPAMetadataBuilderV1

58 +

59 + self.attn_metadata_builders[0].reorder_batch(self.input_batch,

60 + scheduler_output)

61 +

62 def _postprocess_tenosrs(self) -> None:

63 # Note: replace device tensors with cpu tensors

64 def replace_tensor(obj: Any, cpu_attr_name: str,

vllm/v1/worker/gpu_model_runner.py

+34 -15

@@ -49,7 +49,8 @@ from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend

49 from vllm.v1.attention.backends.utils import (

50 AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,

51 make_kv_sharing_fast_prefill_attention_metadata,

52 - make_local_attention_virtual_batches)

52 + make_local_attention_virtual_batches,

53 + reorder_batch_to_split_decodes_and_prefills)

54 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

55 from vllm.v1.kv_cache_interface import (AttentionSpec,

56 ChunkedLocalAttentionSpec,

@@ -329,6 +330,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

330 self.kv_sharing_fast_prefill_logits_indices = torch.zeros(

331 self.max_num_tokens, dtype=torch.int32, device=self.device)

332

333 + self.reorder_batch_threshold: Optional[int] = None

334 +

335 def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:

336 """

337 Update the order of requests in the batch based on the attention

@@ -347,20 +350,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

350 if len(self.kv_cache_config.kv_cache_groups) == 0:

351 return

352

350 - self.attn_metadata_builders[0].reorder_batch(self.input_batch,

351 - scheduler_output)

352 -

353 - # For models with multiple KV cache groups, the groups should agree on

354 - # the same order of requests. We ensure this by only allowing the first

355 - # group to reorder the batch and asserting that all other groups do not

356 - # reorder the batch.

357 - # TODO(tdoublep): make this more flexible so that any group can

358 - # re-order the batch (not only the first).

359 - # TODO(tdoublep): verify this during engine init instead of at runtime

360 - for i in range(1, len(self.kv_cache_config.kv_cache_groups)):

361 - batch_reordered = self.attn_metadata_builders[i].reorder_batch(

362 - self.input_batch, scheduler_output)

363 - assert not batch_reordered

353 + if self.reorder_batch_threshold is not None:

354 + reorder_batch_to_split_decodes_and_prefills(

355 + self.input_batch,

356 + scheduler_output,

357 + decode_threshold=self.reorder_batch_threshold)

358

359 # Note: used for model runner override.

360 def _init_device_properties(self) -> None:

@@ -2654,6 +2648,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2648 self.attn_backends.append(attn_backend_i)

2649 self.attn_metadata_builders.append(attn_metadata_builder_i)

2650

2651 + # Calculate reorder batch threshold (if neeeded)

2652 + self.calculate_reorder_batch_threshold()

2653 +

2654 if len(self.attn_backends) > 0:

2655 return

2656

@@ -2688,6 +2685,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2685 self.attn_metadata_builders.append(attn_metadata_builder)

2686 self.is_encoder_only_model = True

2687

2688 + def calculate_reorder_batch_threshold(self) -> None:

2689 + """

2690 + Check that if any backends reorder batches; that the reordering

2691 + is compatible (e.g., decode threshold is the same)

2692 + """

2693 + for attn_metadata_builder_i in self.attn_metadata_builders:

2694 + # check that if any backends reorder batches; that the reordering

2695 + # is compatible (e.g., decode threshold is the same)

2696 + reorder_batch_threshold_i = (

2697 + attn_metadata_builder_i.reorder_batch_threshold)

2698 + if reorder_batch_threshold_i is not None:

2699 + if self.reorder_batch_threshold is not None:

2700 + if reorder_batch_threshold_i != \

2701 + self.reorder_batch_threshold:

2702 + raise ValueError(

2703 + f"Attention backend reorders decodes with "

2704 + f"threshold {reorder_batch_threshold_i} but other "

2705 + f"backend uses threshold "

2706 + f"{self.reorder_batch_threshold}")

2707 + else:

2708 + self.reorder_batch_threshold = reorder_batch_threshold_i

2709 +

2710 def may_reinitialize_input_batch(self,

2711 kv_cache_config: KVCacheConfig) -> None:

2712 """