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 13, 22:00 UTC
[Hybrid] Pass kernel block size to builders (#27753)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Nov 3, 2025 at 06:48 UTC
18961c5ea62976efc50525b72e40337993c5e4f9
4 files changed
+62
-27
vllm/v1/attention/backends/flash_attn.py
+5
-1
@@ -62,7 +62,11 @@ class FlashAttentionBackend(AttentionBackend):
62
63
@staticmethod
64
def get_supported_kernel_block_size() -> list[int | MultipleOf]:
65
- return [MultipleOf(16)]
65
+ # NOTE(tdoublep): while in principle, FA supports
66
+ # MultipleOf(16), these are the block sizes that do not
67
+ # suffer from the NaN propagation problem described here:
68
+ # https://github.com/Dao-AILab/flash-attention/issues/1974
69
+ return [16, 32, 64]
70
71
@classmethod
72
def validate_head_size(cls, head_size: int) -> None:
vllm/v1/kv_cache_interface.py
+7
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import copy
5
-from dataclasses import dataclass, fields
5
+from dataclasses import dataclass, fields, replace
6
from math import prod
7
8
import torch
@@ -44,6 +44,12 @@ class KVCacheSpec:
44
"""
45
raise NotImplementedError
46
47
+ def copy_with_new_block_size(self, block_size: int) -> Self:
48
+ """
49
+ Create a new KVCacheSpec from self but replacing the block size.
50
+ """
51
+ return replace(self, block_size=block_size)
52
+
53
@classmethod
54
def merge(cls, specs: list[Self]) -> Self:
55
"""
vllm/v1/worker/gpu_model_runner.py
+25
-6
@@ -4039,16 +4039,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4039
) -> list[AttentionGroup]:
4040
attn_groups: list[AttentionGroup] = []
4041
for (attn_backend, kv_cache_spec), layer_names in attn_backends_map.items():
4042
- attn_group = AttentionGroup.create_with_metadata_builders(
4042
+ attn_group = AttentionGroup(
4043
attn_backend,
4044
layer_names,
4045
kv_cache_spec,
4046
- self.vllm_config,
4047
- self.device,
4046
kv_cache_group_id,
4049
- num_metadata_builders=1
4050
- if not self.parallel_config.enable_dbo
4051
- else 2,
4047
)
4048
4049
attn_groups.append(attn_group)
@@ -4067,7 +4062,27 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4062
for i, attn_backend_map in enumerate(attention_backend_maps):
4063
self.attn_groups.append(create_attn_groups(attn_backend_map, i))
4064
4065
+ def initialize_metadata_builders(
4066
+ self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
4067
+ ) -> None:
4068
+ """
4069
+ Create the metadata builders for all KV cache groups and attn groups.
4070
+ """
4071
+ for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)):
4072
+ for attn_group in self.attn_groups[kv_cache_group_id]:
4073
+ attn_group.create_metadata_builders(
4074
+ self.vllm_config,
4075
+ self.device,
4076
+ kernel_block_sizes[kv_cache_group_id]
4077
+ if kv_cache_group_id < len(kernel_block_sizes)
4078
+ else None,
4079
+ num_metadata_builders=1
4080
+ if not self.parallel_config.enable_dbo
4081
+ else 2,
4082
+ )
4083
# Calculate reorder batch threshold (if needed)
4084
+ # Note (tdoublep): do this *after* constructing builders,
4085
+ # because some of them change the threshold at init time.
4086
self.calculate_reorder_batch_threshold()
4087
4088
def _check_and_update_cudagraph_mode(
@@ -4633,6 +4648,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4648
# kernel_block_size 64 and split the 256-token-block to 4 blocks with 64
4649
# tokens each.
4650
kernel_block_sizes = self._prepare_kernel_block_sizes(kv_cache_config)
4651
+
4652
+ # create metadata builders
4653
+ self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)
4654
+
4655
# Reinitialize need to after initialize_attn_backend
4656
self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
4657
kv_caches = self.initialize_kv_cache_tensors(
vllm/v1/worker/utils.py
+25
-19
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections import defaultdict
4
-from dataclasses import dataclass
4
+from dataclasses import dataclass, field
5
from typing import TYPE_CHECKING
6
7
import torch
@@ -134,31 +134,37 @@ class MultiModalBudget:
134
@dataclass
135
class AttentionGroup:
136
backend: type[AttentionBackend]
137
- # When ubatching is enabled we will have a metadata builder for each ubatch
138
- # so that if they use internal persistant buffers for cudagraphs, and they
139
- # won't have to worry about conflicting with the other ubatches.
140
- metadata_builders: list[AttentionMetadataBuilder]
137
layer_names: list[str]
138
kv_cache_spec: KVCacheSpec
139
kv_cache_group_id: int
140
+ # When ubatching is enabled we will have a metadata builder for each ubatch
141
+ # so that if they use internal persistant buffers for cudagraphs, and they
142
+ # won't have to worry about conflicting with the other ubatches.
143
+ metadata_builders: list[AttentionMetadataBuilder] = field(
144
+ default_factory=lambda: []
145
+ )
146
145
- @staticmethod
146
- def create_with_metadata_builders(
147
- backend: type[AttentionBackend],
148
- layer_names: list[str],
149
- kv_cache_spec: KVCacheSpec,
150
- vllm_config: VllmConfig,
151
- device: torch.device,
152
- kv_cache_group_id: int,
147
+ def create_metadata_builders(
148
+ self,
149
+ vllm_config,
150
+ device,
151
+ kernel_block_size: int | None,
152
num_metadata_builders: int = 1,
154
- ) -> "AttentionGroup":
155
- metadata_builders = [
156
- backend.get_builder_cls()(kv_cache_spec, layer_names, vllm_config, device)
153
+ ):
154
+ kv_cache_spec_builder = (
155
+ self.kv_cache_spec.copy_with_new_block_size(kernel_block_size)
156
+ if kernel_block_size is not None
157
+ else self.kv_cache_spec
158
+ )
159
+ self.metadata_builders = [
160
+ self.backend.get_builder_cls()(
161
+ kv_cache_spec_builder,
162
+ self.layer_names,
163
+ vllm_config,
164
+ device,
165
+ )
166
for _ in range(num_metadata_builders)
167
]
159
- return AttentionGroup(
160
- backend, metadata_builders, layer_names, kv_cache_spec, kv_cache_group_id
161
- )
168
169
def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder:
170
assert len(self.metadata_builders) > ubatch_id