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 7, 15:15 UTC
[Misc] Move `GPUModelRunner.prepare_kernel_block_sizes` to utils (#35400)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Feb 27, 2026 at 04:42 UTC
cabdaa761975179afcc90dc6965430bce033e098
3 files changed
+135
-123
tests/v1/worker/test_gpu_model_runner.py
+4
-4
@@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import (
38
from vllm.v1.sample.metadata import SamplingMetadata
39
from vllm.v1.worker.gpu_input_batch import InputBatch
40
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
41
-from vllm.v1.worker.utils import AttentionGroup
41
+from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
42
43
BLOCK_SIZE = 16
44
NUM_BLOCKS = 10
@@ -209,7 +209,7 @@ def test_select_common_block_size_prefers_manager_block_size():
209
AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),
210
]
211
212
- selected_size = GPUModelRunner.select_common_block_size(128, attn_groups)
212
+ selected_size = select_common_block_size(128, attn_groups)
213
assert selected_size == 128
214
215
@@ -221,7 +221,7 @@ def test_select_common_block_size_uses_largest_shared_int():
221
AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),
222
]
223
224
- selected_size = GPUModelRunner.select_common_block_size(256, attn_groups)
224
+ selected_size = select_common_block_size(256, attn_groups)
225
assert selected_size == 64
226
227
@@ -234,7 +234,7 @@ def test_select_common_block_size_no_valid_option():
234
]
235
236
with pytest.raises(ValueError):
237
- GPUModelRunner.select_common_block_size(48, attn_groups)
237
+ select_common_block_size(48, attn_groups)
238
239
240
def test_update_states_new_request(model_runner, dist_init):
vllm/v1/worker/gpu_model_runner.py
+4
-117
@@ -115,7 +115,6 @@ from vllm.v1.attention.backend import (
115
AttentionMetadataBuilder,
116
AttentionType,
117
CommonAttentionMetadata,
118
- MultipleOf,
118
)
119
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
120
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
@@ -189,6 +188,7 @@ from .utils import (
188
AttentionGroup,
189
add_kv_sharing_layers_to_kv_cache_groups,
190
bind_kv_cache,
191
+ prepare_kernel_block_sizes,
192
sanity_check_mm_encoder_outputs,
193
)
194
@@ -5678,78 +5678,6 @@ class GPUModelRunner(
5678
return
5679
self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment]
5680
5681
- @staticmethod
5682
- def select_common_block_size(
5683
- kv_manager_block_size: int, attn_groups: list[AttentionGroup]
5684
- ) -> int:
5685
- """
5686
- Select a block size that is supported by all backends and is a factor of
5687
- kv_manager_block_size.
5688
-
5689
- If kv_manager_block_size is supported by all backends, return it directly.
5690
- Otherwise, return the max supported size.
5691
-
5692
- Args:
5693
- kv_manager_block_size: Block size of KV cache
5694
- attn_groups: List of attention groups
5695
-
5696
- Returns:
5697
- The selected block size
5698
-
5699
- Raises:
5700
- ValueError: If no valid block size found
5701
- """
5702
-
5703
- def block_size_is_supported(
5704
- backends: list[type[AttentionBackend]], block_size: int
5705
- ) -> bool:
5706
- """
5707
- Check if the block size is supported by all backends.
5708
- """
5709
- for backend in backends:
5710
- is_supported = False
5711
- for supported_size in backend.get_supported_kernel_block_sizes():
5712
- if isinstance(supported_size, int):
5713
- if block_size == supported_size:
5714
- is_supported = True
5715
- elif isinstance(supported_size, MultipleOf):
5716
- if block_size % supported_size.base == 0:
5717
- is_supported = True
5718
- else:
5719
- raise ValueError(f"Unknown supported size: {supported_size}")
5720
- if not is_supported:
5721
- return False
5722
- return True
5723
-
5724
- backends = [group.backend for group in attn_groups]
5725
-
5726
- # Case 1: if the block_size of kv cache manager is supported by all backends,
5727
- # return it directly
5728
- if block_size_is_supported(backends, kv_manager_block_size):
5729
- return kv_manager_block_size
5730
-
5731
- # Case 2: otherwise, the block_size must be an `int`-format supported size of
5732
- # at least one backend. Iterate over all `int`-format supported sizes in
5733
- # descending order and return the first one that is supported by all backends.
5734
- # Simple proof:
5735
- # If the supported size b is in MultipleOf(x_i) format for all attention
5736
- # backends i, and b a factor of kv_manager_block_size, then
5737
- # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will
5738
- # return kv_manager_block_size in case 1.
5739
- all_int_supported_sizes = set(
5740
- supported_size
5741
- for backend in backends
5742
- for supported_size in backend.get_supported_kernel_block_sizes()
5743
- if isinstance(supported_size, int)
5744
- )
5745
-
5746
- for supported_size in sorted(all_int_supported_sizes, reverse=True):
5747
- if kv_manager_block_size % supported_size != 0:
5748
- continue
5749
- if block_size_is_supported(backends, supported_size):
5750
- return supported_size
5751
- raise ValueError(f"No common block size for {kv_manager_block_size}. ")
5752
-
5681
def may_reinitialize_input_batch(
5682
self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
5683
) -> None:
@@ -5846,49 +5774,6 @@ class GPUModelRunner(
5774
for attn_groups in self.attn_groups:
5775
yield from attn_groups
5776
5849
- def _prepare_kernel_block_sizes(self, kv_cache_config: KVCacheConfig) -> list[int]:
5850
- """
5851
- Generate kernel_block_sizes that matches each block_size.
5852
-
5853
- For attention backends that support virtual block splitting,
5854
- use the supported block sizes from the backend.
5855
- For other backends (like Mamba), use the same block size (no splitting).
5856
-
5857
- Args:
5858
- kv_cache_config: The KV cache configuration.
5859
-
5860
- Returns:
5861
- list[int]: List of kernel block sizes for each cache group.
5862
- """
5863
- kernel_block_sizes = []
5864
- for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
5865
- kv_cache_spec = kv_cache_group.kv_cache_spec
5866
- if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
5867
- # All layers in the UniformTypeKVCacheSpecs have the same type,
5868
- # Pick an arbitrary one to dispatch.
5869
- kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
5870
- if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
5871
- continue
5872
- elif isinstance(kv_cache_spec, AttentionSpec):
5873
- # This is an attention backend that supports virtual
5874
- # block splitting. Get the supported block sizes from
5875
- # all backends in the group.
5876
- attn_groups = self.attn_groups[kv_cache_gid]
5877
- kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
5878
- selected_kernel_size = self.select_common_block_size(
5879
- kv_manager_block_size, attn_groups
5880
- )
5881
- kernel_block_sizes.append(selected_kernel_size)
5882
- elif isinstance(kv_cache_spec, MambaSpec):
5883
- # This is likely Mamba or other non-attention cache,
5884
- # no splitting.
5885
- kernel_block_sizes.append(kv_cache_spec.block_size)
5886
- else:
5887
- raise NotImplementedError(
5888
- f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"
5889
- )
5890
- return kernel_block_sizes
5891
-
5777
def _reshape_kv_cache_tensors(
5778
self,
5779
kv_cache_config: KVCacheConfig,
@@ -6120,7 +6005,9 @@ class GPUModelRunner(
6005
# backends for that group only supports block_size 64, we will return
6006
# kernel_block_size 64 and split the 256-token-block to 4 blocks with 64
6007
# tokens each.
6123
- kernel_block_sizes = self._prepare_kernel_block_sizes(kv_cache_config)
6008
+ kernel_block_sizes = prepare_kernel_block_sizes(
6009
+ kv_cache_config, self.attn_groups
6010
+ )
6011
6012
# create metadata builders
6013
self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)
vllm/v1/worker/utils.py
+127
-2
@@ -13,8 +13,20 @@ from vllm.model_executor.models.interfaces import MultiModalEmbeddings
13
from vllm.model_executor.models.utils import extract_layer_index
14
from vllm.platforms import current_platform
15
from vllm.utils.mem_utils import MemorySnapshot, format_gib
16
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadataBuilder
17
-from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
16
+from vllm.v1.attention.backend import (
17
+ AttentionBackend,
18
+ AttentionMetadataBuilder,
19
+ MultipleOf,
20
+)
21
+from vllm.v1.kv_cache_interface import (
22
+ AttentionSpec,
23
+ EncoderOnlyAttentionSpec,
24
+ KVCacheConfig,
25
+ KVCacheGroupSpec,
26
+ KVCacheSpec,
27
+ MambaSpec,
28
+ UniformTypeKVCacheSpecs,
29
+)
30
31
logger = init_logger(__name__)
32
@@ -59,6 +71,119 @@ class AttentionGroup:
71
return self.metadata_builders[ubatch_id]
72
73
74
+def select_common_block_size(
75
+ kv_manager_block_size: int, attn_groups: list[AttentionGroup]
76
+) -> int:
77
+ """
78
+ Select a block size that is supported by all backends and is a factor of
79
+ kv_manager_block_size.
80
+
81
+ If kv_manager_block_size is supported by all backends, return it directly.
82
+ Otherwise, return the max supported size.
83
+
84
+ Args:
85
+ kv_manager_block_size: Block size of KV cache.
86
+ attn_groups: List of attention groups.
87
+
88
+ Returns:
89
+ The selected block size.
90
+
91
+ Raises:
92
+ ValueError: If no valid block size found.
93
+ """
94
+
95
+ def block_size_is_supported(
96
+ backends: list[type[AttentionBackend]], block_size: int
97
+ ) -> bool:
98
+ """Check if the block size is supported by all backends."""
99
+ for backend in backends:
100
+ is_supported = False
101
+ for supported_size in backend.get_supported_kernel_block_sizes():
102
+ if isinstance(supported_size, int):
103
+ if block_size == supported_size:
104
+ is_supported = True
105
+ elif isinstance(supported_size, MultipleOf):
106
+ if block_size % supported_size.base == 0:
107
+ is_supported = True
108
+ else:
109
+ raise ValueError(f"Unknown supported size: {supported_size}")
110
+ if not is_supported:
111
+ return False
112
+ return True
113
+
114
+ backends = [group.backend for group in attn_groups]
115
+
116
+ # Case 1: if the block_size of kv cache manager is supported by all backends,
117
+ # return it directly.
118
+ if block_size_is_supported(backends, kv_manager_block_size):
119
+ return kv_manager_block_size
120
+
121
+ # Case 2: otherwise, the block_size must be an `int`-format supported size of
122
+ # at least one backend. Iterate over all `int`-format supported sizes in
123
+ # descending order and return the first one that is supported by all backends.
124
+ # Simple proof:
125
+ # If the supported size b is in MultipleOf(x_i) format for all attention
126
+ # backends i, and b a factor of kv_manager_block_size, then
127
+ # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will
128
+ # return kv_manager_block_size in case 1.
129
+ all_int_supported_sizes = set(
130
+ supported_size
131
+ for backend in backends
132
+ for supported_size in backend.get_supported_kernel_block_sizes()
133
+ if isinstance(supported_size, int)
134
+ )
135
+
136
+ for supported_size in sorted(all_int_supported_sizes, reverse=True):
137
+ if kv_manager_block_size % supported_size != 0:
138
+ continue
139
+ if block_size_is_supported(backends, supported_size):
140
+ return supported_size
141
+ raise ValueError(f"No common block size for {kv_manager_block_size}. ")
142
+
143
+
144
+def prepare_kernel_block_sizes(
145
+ kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]]
146
+) -> list[int]:
147
+ """
148
+ Generate kernel_block_sizes that matches each block_size.
149
+
150
+ For attention backends that support virtual block splitting,
151
+ use the supported block sizes from the backend.
152
+ For other backends (like Mamba), use the same block size (no splitting).
153
+
154
+ Args:
155
+ kv_cache_config: The KV cache configuration.
156
+ attn_groups: Attention groups indexed by KV cache group id.
157
+
158
+ Returns:
159
+ List of kernel block sizes for each cache group.
160
+ """
161
+ kernel_block_sizes = []
162
+ for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
163
+ kv_cache_spec = kv_cache_group.kv_cache_spec
164
+ if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
165
+ # All layers in the UniformTypeKVCacheSpecs have the same type,
166
+ # pick an arbitrary one to dispatch.
167
+ kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
168
+ if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
169
+ continue
170
+ if isinstance(kv_cache_spec, AttentionSpec):
171
+ # This is an attention backend that supports virtual block splitting.
172
+ kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
173
+ selected_kernel_size = select_common_block_size(
174
+ kv_manager_block_size, attn_groups[kv_cache_gid]
175
+ )
176
+ kernel_block_sizes.append(selected_kernel_size)
177
+ elif isinstance(kv_cache_spec, MambaSpec):
178
+ # This is likely Mamba or other non-attention cache, no splitting.
179
+ kernel_block_sizes.append(kv_cache_spec.block_size)
180
+ else:
181
+ raise NotImplementedError(
182
+ f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"
183
+ )
184
+ return kernel_block_sizes
185
+
186
+
187
def sanity_check_mm_encoder_outputs(
188
mm_embeddings: MultiModalEmbeddings,
189
expected_num_items: int,