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 8, 01:15 UTC
[Bugfix] Fix Basic Models Test (#34818)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Matthew Bonanni committed
Feb 19, 2026 at 17:49 UTC
662205d34eb1bb42228768d7a69a1ac4abf38c89
14 files changed
+174
-220
tests/models/multimodal/processing/test_tensor_schema.py
+4
-1
@@ -13,6 +13,7 @@ import torch.nn as nn
13
from PIL import Image
14
15
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
16
+from vllm.config.cache import CacheConfig
17
from vllm.config.multimodal import (
18
AudioDummyOptions,
19
BaseDummyOptions,
@@ -131,7 +132,9 @@ def initialize_dummy_model(
132
):
133
temp_file = tempfile.mkstemp()[1]
134
current_device = torch.get_default_device()
134
- vllm_config = VllmConfig(model_config=model_config)
135
+ vllm_config = VllmConfig(
136
+ model_config=model_config, cache_config=CacheConfig(block_size=16)
137
+ )
138
with set_current_vllm_config(vllm_config=vllm_config):
139
init_distributed_environment(
140
world_size=1,
tests/models/utils.py
+7
-2
@@ -457,6 +457,9 @@ def dummy_hf_overrides(
457
# Kimi uses `num_expert_group` instead of `n_group`.
458
if n_group is None:
459
n_group = getattr(text_config, "num_expert_group", None)
460
+ # InternS1Pro uses `router_n_groups` instead of `n_group`.
461
+ if n_group is None:
462
+ n_group = getattr(text_config, "router_n_groups", None)
463
num_experts = n_group * 2 if n_group is not None else 2
464
465
# we use three layers for Gemma-3n to check
@@ -486,12 +489,14 @@ def dummy_hf_overrides(
489
# Only set MoE related config when the model has MoE layers.
490
# Otherwise all models detected as MoE by _get_transformers_backend_cls.
491
if model_arch_config.num_experts > 0:
492
+ orig_topk = getattr(text_config, "num_experts_per_tok", 2)
493
+ topk = min(orig_topk, 2)
494
update_dict.update(
495
{
496
"num_experts": num_experts,
492
- "num_experts_per_tok": 2,
497
+ "num_experts_per_tok": topk,
498
# Kimi uses `num_experts_per_token`.
494
- "num_experts_per_token": 2,
499
+ "num_experts_per_token": topk,
500
"num_local_experts": num_experts,
501
# Otherwise there will not be any expert layers
502
"first_k_dense_replace": 0,
tests/v1/spec_decode/test_eagle.py
+1
-1
@@ -78,7 +78,7 @@ def _create_proposer(
78
device = current_platform.device_type
79
vllm_config = VllmConfig(
80
model_config=model_config,
81
- cache_config=CacheConfig(),
81
+ cache_config=CacheConfig(block_size=16),
82
speculative_config=speculative_config,
83
device_config=DeviceConfig(device=device),
84
parallel_config=ParallelConfig(),
vllm/config/cache.py
+2
-2
@@ -41,8 +41,8 @@ class CacheConfig:
41
block_size: SkipValidation[int] = None # type: ignore[assignment]
42
"""Size of a contiguous cache block in number of tokens.
43
44
- This is None until `Platform.check_and_update_config()` sets it based on
45
- the current platform. Always an int by the time the engine starts."""
44
+ This is None until the platform sets it. Always an int by the time
45
+ the engine starts."""
46
gpu_memory_utilization: float = Field(default=0.9, gt=0, le=1)
47
"""The fraction of GPU memory to be used for the model executor, which can
48
range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory
vllm/config/vllm.py
+51
-46
@@ -915,32 +915,6 @@ class VllmConfig:
915
)
916
current_platform.check_and_update_config(self)
917
918
- # If DCP, ensure the block size is right.
919
- if self.parallel_config.decode_context_parallel_size > 1:
920
- if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
921
- self.parallel_config.cp_kv_cache_interleave_size
922
- != self.parallel_config.dcp_kv_cache_interleave_size
923
- ):
924
- self.parallel_config.cp_kv_cache_interleave_size = (
925
- self.parallel_config.dcp_kv_cache_interleave_size
926
- )
927
- logger.warning_once(
928
- "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
929
- "_interleave_size. And dcp-kv-cache-interleave-size will be "
930
- "deprecated when PCP is fully supported."
931
- )
932
- assert (
933
- self.parallel_config.cp_kv_cache_interleave_size
934
- <= self.cache_config.block_size
935
- and self.cache_config.block_size
936
- % self.parallel_config.cp_kv_cache_interleave_size
937
- == 0
938
- ), (
939
- f"Block_size({self.cache_config.block_size}) should be greater "
940
- "than or equal to and divisible by cp_kv_cache_interleave_size "
941
- f"({self.parallel_config.cp_kv_cache_interleave_size})."
942
- )
943
-
918
# Do this after all the updates to compilation_config.mode
919
effective_dp_size = (
920
self.parallel_config.data_parallel_size
@@ -1108,26 +1082,6 @@ class VllmConfig:
1082
# Default to enable HMA if not explicitly disabled by user or logic above.
1083
self.scheduler_config.disable_hybrid_kv_cache_manager = False
1084
1111
- if self.cache_config.mamba_cache_mode == "align":
1112
- assert (
1113
- self.cache_config.block_size
1114
- <= self.scheduler_config.max_num_batched_tokens
1115
- ), (
1116
- "In Mamba cache align mode, block_size "
1117
- f"({self.cache_config.block_size}) must be <= "
1118
- "max_num_batched_tokens "
1119
- f"({self.scheduler_config.max_num_batched_tokens})."
1120
- )
1121
- if self.scheduler_config.long_prefill_token_threshold > 0:
1122
- assert (
1123
- self.scheduler_config.long_prefill_token_threshold
1124
- >= self.cache_config.block_size
1125
- )
1126
- assert not self.scheduler_config.disable_chunked_mm_input, (
1127
- "Chunked MM input is required because we need the flexibility to "
1128
- "schedule a multiple of block_size tokens even if they are in the "
1129
- "middle of a mm input"
1130
- )
1085
if self.compilation_config.debug_dump_path:
1086
self.compilation_config.debug_dump_path = (
1087
self.compilation_config.debug_dump_path.absolute().expanduser()
@@ -1488,6 +1442,57 @@ class VllmConfig:
1442
f"compilation_config={self.compilation_config!r}"
1443
)
1444
1445
+ def validate_block_size(self) -> None:
1446
+ """Validate block_size against DCP and mamba constraints.
1447
+
1448
+ Called after Platform.update_block_size_for_backend() has
1449
+ finalised block_size, so that the checks see the real value
1450
+ rather than the initial None sentinel.
1451
+ """
1452
+ block_size = self.cache_config.block_size
1453
+ assert block_size is not None, (
1454
+ "validate_block_size called before block_size was set"
1455
+ )
1456
+
1457
+ # DCP interleave-size compatibility
1458
+ if self.parallel_config.decode_context_parallel_size > 1:
1459
+ if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
1460
+ self.parallel_config.cp_kv_cache_interleave_size
1461
+ != self.parallel_config.dcp_kv_cache_interleave_size
1462
+ ):
1463
+ self.parallel_config.cp_kv_cache_interleave_size = (
1464
+ self.parallel_config.dcp_kv_cache_interleave_size
1465
+ )
1466
+ logger.warning_once(
1467
+ "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
1468
+ "_interleave_size. And dcp-kv-cache-interleave-size will be "
1469
+ "deprecated when PCP is fully supported."
1470
+ )
1471
+ assert (
1472
+ self.parallel_config.cp_kv_cache_interleave_size <= block_size
1473
+ and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0
1474
+ ), (
1475
+ f"Block_size({block_size}) should be greater "
1476
+ "than or equal to and divisible by cp_kv_cache_interleave_size "
1477
+ f"({self.parallel_config.cp_kv_cache_interleave_size})."
1478
+ )
1479
+
1480
+ # Mamba cache align-mode constraints
1481
+ if self.cache_config.mamba_cache_mode == "align":
1482
+ assert block_size <= self.scheduler_config.max_num_batched_tokens, (
1483
+ "In Mamba cache align mode, block_size "
1484
+ f"({block_size}) must be <= "
1485
+ "max_num_batched_tokens "
1486
+ f"({self.scheduler_config.max_num_batched_tokens})."
1487
+ )
1488
+ if self.scheduler_config.long_prefill_token_threshold > 0:
1489
+ assert self.scheduler_config.long_prefill_token_threshold >= block_size
1490
+ assert not self.scheduler_config.disable_chunked_mm_input, (
1491
+ "Chunked MM input is required because we need the flexibility "
1492
+ "to schedule a multiple of block_size tokens even if they are "
1493
+ "in the middle of a mm input"
1494
+ )
1495
+
1496
@model_validator(mode="after")
1497
def validate_mamba_block_size(self) -> "VllmConfig":
1498
if self.model_config is None:
vllm/model_executor/layers/attention/chunked_local_attention.py
+6
-5
@@ -30,9 +30,8 @@ from vllm.v1.kv_cache_interface import (
30
def create_chunked_local_attention_backend(
31
underlying_attn_backend: AttentionBackend,
32
attention_chunk_size: int,
33
- block_size: int,
33
) -> type[AttentionBackend]:
35
- prefix = f"ChunkedLocalAttention_{attention_chunk_size}_{block_size}_"
34
+ prefix = f"ChunkedLocalAttention_{attention_chunk_size}_"
35
36
underlying_builder = underlying_attn_backend.get_builder_cls()
37
assert issubclass(underlying_builder, AttentionMetadataBuilder)
@@ -55,7 +54,9 @@ def create_chunked_local_attention_backend(
54
fast_build: bool = False,
55
):
56
cm, make_virtual_batches_block_table = make_local_attention_virtual_batches(
58
- attention_chunk_size, common_attn_metadata, block_size
57
+ attention_chunk_size,
58
+ common_attn_metadata,
59
+ self.kv_cache_spec.block_size,
60
)
61
metadata = super().build(common_prefix_len, cm, fast_build)
62
metadata.make_virtual_batches_block_table = make_virtual_batches_block_table
@@ -97,13 +98,13 @@ class ChunkedLocalAttention(Attention):
98
block_size = cache_config.block_size
99
else:
100
kv_cache_dtype = "auto"
100
- block_size = 16
101
+ block_size = None
102
103
underlying_attn_backend = get_attn_backend(
104
head_size, dtype, kv_cache_dtype, block_size
105
)
106
attn_backend = create_chunked_local_attention_backend(
106
- underlying_attn_backend, attention_chunk_size, block_size
107
+ underlying_attn_backend, attention_chunk_size
108
)
109
110
super().__init__(
vllm/model_executor/layers/attention/mla_attention.py
+12
-5
@@ -407,17 +407,24 @@ class MLAAttention(nn.Module, AttentionLayerBase):
407
)
408
409
# Attributes for forward_impl method
410
- self.chunked_prefill_workspace_size = (
411
- MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
412
- get_current_vllm_config()
413
- )
414
- )
410
+ self._vllm_config = get_current_vllm_config()
411
+ self._chunked_prefill_workspace_size: int | None = None
412
self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(
413
static=True,
414
group_shape=GroupShape.PER_TENSOR,
415
compile_native=True,
416
)
417
418
+ @property
419
+ def chunked_prefill_workspace_size(self) -> int:
420
+ if self._chunked_prefill_workspace_size is None:
421
+ self._chunked_prefill_workspace_size = (
422
+ MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
423
+ self._vllm_config
424
+ )
425
+ )
426
+ return self._chunked_prefill_workspace_size
427
+
428
def forward(
429
self,
430
q: torch.Tensor,
vllm/platforms/cuda.py
+62
-155
@@ -169,21 +169,6 @@ class CudaPlatformBase(Platform):
169
if parallel_config.worker_cls == "auto":
170
parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
171
172
- cache_config = vllm_config.cache_config
173
- user_specified_block_size = cache_config.block_size is not None
174
- if not user_specified_block_size:
175
- cache_config.block_size = 16
176
-
177
- # Ensure block_size is compatible with the attention backend.
178
- # Note: model_config may be None during testing.
179
- # Skip hybrid (attention+mamba) models — their block_size is
180
- # managed by HybridAttentionMambaModelConfig
181
- if model_config is not None and not model_config.is_hybrid:
182
- cls._update_block_size_for_backend(
183
- vllm_config,
184
- user_specified_block_size,
185
- )
186
-
172
scheduler_config = vllm_config.scheduler_config
173
# Note: model_config may be None during testing
174
if (
@@ -199,148 +184,47 @@ class CudaPlatformBase(Platform):
184
scheduler_config.disable_chunked_mm_input = True
185
186
@classmethod
202
- def _update_block_size_for_backend(
203
- cls,
204
- vllm_config: "VllmConfig",
205
- user_specified_block_size: bool,
206
- ) -> None:
207
- """Ensure block_size is compatible with the attention backend.
208
-
209
- If the user specified --block-size, the selector validates/filters
210
- backends by that block size (raising on incompatibility). Otherwise,
211
- the backend is selected unconstrained and block_size is set to the
212
- backend's preferred value.
213
- """
214
- from vllm.config.vllm import set_current_vllm_config
215
- from vllm.v1.attention.selector import AttentionSelectorConfig
216
-
217
- model_config = vllm_config.model_config
187
+ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
188
cache_config = vllm_config.cache_config
189
+ if cache_config.block_size is not None:
190
+ # User specified --block-size; keep it.
191
+ return
192
220
- device_capability = cls.get_device_capability()
221
- if device_capability is None:
193
+ model_config = vllm_config.model_config
194
+ # model_config may be None during testing.
195
+ # Skip hybrid models — their block_size is managed by
196
+ # HybridAttentionMambaModelConfig.
197
+ if model_config is None or model_config.is_hybrid:
198
+ cache_config.block_size = 16
199
return
200
224
- use_mla = model_config.use_mla
225
- attn_selector_config = AttentionSelectorConfig(
226
- head_size=model_config.get_head_size(),
227
- dtype=model_config.dtype, # type: ignore[arg-type]
228
- kv_cache_dtype=cache_config.cache_dtype,
229
- block_size=cache_config.block_size if user_specified_block_size else None,
230
- use_mla=use_mla,
231
- has_sink=False,
232
- use_sparse=use_mla and hasattr(model_config.hf_config, "index_topk"),
233
- use_mm_prefix=model_config.is_mm_prefix_lm,
201
+ from vllm.config.vllm import (
202
+ get_layers_from_vllm_config,
203
+ set_current_vllm_config,
204
+ )
205
+ from vllm.model_executor.layers.attention_layer_base import (
206
+ AttentionLayerBase,
207
)
208
236
- user_specified_backend = vllm_config.attention_config.backend
237
- num_heads = model_config.get_num_attention_heads(
238
- vllm_config.parallel_config,
209
+ attn_layers = get_layers_from_vllm_config(
210
+ vllm_config,
211
+ AttentionLayerBase,
212
)
213
+ if not attn_layers:
214
+ cache_config.block_size = 16
215
+ return
216
+
217
+ first_layer = next(iter(attn_layers.values()))
218
+ backend_cls = first_layer.get_attn_backend()
219
with set_current_vllm_config(vllm_config):
241
- chosen_backend = cls.select_attention_backend(
242
- selected_backend=user_specified_backend,
243
- attn_selector_config=attn_selector_config,
244
- device_capability=device_capability,
245
- # Don't raise here — we produce better errors below.
246
- raise_on_invalid=False,
247
- num_heads=num_heads,
220
+ preferred = backend_cls.get_preferred_block_size(16)
221
+ if preferred != 16:
222
+ logger.info(
223
+ "Setting kv cache block size to %d for %s backend.",
224
+ preferred,
225
+ backend_cls.get_name(),
226
)
249
-
250
- # If the user's --block-size forced a non-optimal backend,
251
- # warn them. Only relevant when the user didn't also specify
252
- # --attention-backend (in which case the choice is explicit).
253
- if (
254
- chosen_backend is not None
255
- and user_specified_block_size
256
- and user_specified_backend is None
257
- ):
258
- optimal = cls.select_attention_backend(
259
- selected_backend=None,
260
- attn_selector_config=attn_selector_config._replace(
261
- block_size=None,
262
- ),
263
- device_capability=device_capability,
264
- raise_on_invalid=False,
265
- num_heads=num_heads,
266
- )
267
- if optimal is not None and optimal != chosen_backend:
268
- logger.warning(
269
- "--block-size %d is not supported by the preferred "
270
- "%s backend. Using %s instead, which may result "
271
- "in reduced performance. Consider removing "
272
- "--block-size to auto-select the optimal "
273
- "block size.",
274
- cache_config.block_size,
275
- optimal.name,
276
- chosen_backend.name,
277
- )
278
-
279
- if chosen_backend is not None:
280
- if user_specified_block_size:
281
- # User's block_size is compatible with the chosen
282
- # backend.
283
- return
284
- # User didn't specify --block-size, so auto-select the
285
- # preferred block size for the chosen backend.
286
- try:
287
- backend_class = chosen_backend.get_class()
288
- except ImportError:
289
- return # Will fail later with a better error
290
- preferred = backend_class.get_preferred_block_size(
291
- cache_config.block_size,
292
- )
293
- if cache_config.block_size != preferred:
294
- logger.info(
295
- "Setting kv cache block size to %d for %s backend.",
296
- preferred,
297
- chosen_backend.name,
298
- )
299
- cache_config.block_size = preferred
300
- return
301
-
302
- # No valid backend found. If the user didn't constrain the
303
- # selection, defer the error to get_attn_backend_cls where
304
- # the full config (including per-layer settings) is
305
- # available.
306
- if not user_specified_block_size:
307
- return
308
-
309
- if user_specified_backend is not None:
310
- # User specified --block-size and --attention-backend
311
- # and they are incompatible.
312
- try:
313
- backend_class = user_specified_backend.get_class()
314
- supported = backend_class.get_supported_kernel_block_sizes()
315
- except ImportError:
316
- supported = None
317
- raise ValueError(
318
- f"User-specified --block-size "
319
- f"{cache_config.block_size} is incompatible with "
320
- f"the specified --attention-backend "
321
- f"{user_specified_backend.name} (supported kernel "
322
- f"block sizes: {supported}). Either remove "
323
- f"--block-size to auto-select, or choose a "
324
- f"compatible value."
325
- )
326
- else:
327
- # User specified --block-size but no backend supports
328
- # it.
329
- _, invalid_reasons = cls.get_valid_backends(
330
- device_capability=device_capability,
331
- attn_selector_config=attn_selector_config,
332
- num_heads=num_heads,
333
- )
334
- reasons_str = ", ".join(
335
- f"{b.name}: [{', '.join(r)}]" for b, r in invalid_reasons.items()
336
- )
337
- raise ValueError(
338
- f"No valid attention backend found for "
339
- f"--block-size {cache_config.block_size}. "
340
- f"Reasons: {{{reasons_str}}}. Either remove "
341
- f"--block-size to auto-select, or choose a "
342
- f"compatible value."
343
- )
227
+ cache_config.block_size = preferred
228
229
@classmethod
230
def get_current_memory_usage(
@@ -358,10 +242,10 @@ class CudaPlatformBase(Platform):
242
num_heads: int | None = None,
243
) -> tuple[
244
list[tuple["AttentionBackendEnum", int]],
361
- dict["AttentionBackendEnum", list[str]],
245
+ dict["AttentionBackendEnum", tuple[int, list[str]]],
246
]:
247
valid_backends_priorities = []
364
- invalid_reasons = {}
248
+ invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
249
250
backend_priorities = _get_backend_priorities(
251
attn_selector_config.use_mla,
@@ -378,7 +262,7 @@ class CudaPlatformBase(Platform):
262
except ImportError:
263
invalid_reasons_i = ["ImportError"]
264
if invalid_reasons_i:
381
- invalid_reasons[backend] = invalid_reasons_i
265
+ invalid_reasons[backend] = (priority, invalid_reasons_i)
266
else:
267
valid_backends_priorities.append((backend, priority))
268
@@ -439,7 +323,7 @@ class CudaPlatformBase(Platform):
323
"{"
324
+ ", ".join(
325
f"{backend.name}: [{', '.join(reasons)}]"
442
- for backend, reasons in invalid_reasons.items()
326
+ for backend, (_, reasons) in invalid_reasons.items()
327
)
328
+ "}"
329
)
@@ -452,7 +336,30 @@ class CudaPlatformBase(Platform):
336
337
# Select the one with the highest priority (lowest index).
338
sorted_backends = sorted(valid_backends_priorities, key=lambda x: x[1])
455
- return sorted_backends[0][0]
339
+ chosen_backend, chosen_priority = sorted_backends[0]
340
+
341
+ # If the user specified --block-size (but not --attention-backend),
342
+ # check whether that constraint precluded any higher-priority backends.
343
+ if attn_selector_config.block_size is not None:
344
+ excluded = [
345
+ backend
346
+ for backend, (priority, reasons) in invalid_reasons.items()
347
+ if priority < chosen_priority
348
+ and reasons == ["block_size not supported"]
349
+ ]
350
+ if excluded:
351
+ names = ", ".join(b.name for b in excluded)
352
+ logger.warning(
353
+ "--block-size %d excluded higher-priority backend(s) "
354
+ "%s. Using %s instead, which may result in reduced "
355
+ "performance. Consider removing --block-size to "
356
+ "auto-select the optimal block size.",
357
+ attn_selector_config.block_size,
358
+ names,
359
+ chosen_backend.name,
360
+ )
361
+
362
+ return chosen_backend
363
364
@classmethod
365
def get_attn_backend_cls(
@@ -487,7 +394,7 @@ class CudaPlatformBase(Platform):
394
"{"
395
+ ", ".join(
396
f"{backend.name}: [{', '.join(reasons)}]"
490
- for backend, reasons in invalid_reasons.items()
397
+ for backend, (_, reasons) in invalid_reasons.items()
398
)
399
+ "}"
400
)
@@ -499,7 +406,7 @@ class CudaPlatformBase(Platform):
406
logger.info_once(
407
"Using %s attention backend out of potential backends: %s",
408
chosen_backend.name,
502
- tuple(b[0].name for b in valid_backends_priorities),
409
+ tuple(backend.name for backend, _ in valid_backends_priorities),
410
scope="local",
411
)
412
vllm/platforms/interface.py
+7
@@ -406,6 +406,13 @@ class Platform:
406
"""
407
pass
408
409
+ @classmethod
410
+ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
411
+ """
412
+ Ensure block_size is compatible with the attention backend.
413
+ """
414
+ pass
415
+
416
@classmethod
417
def verify_model_arch(cls, model_arch: str) -> None:
418
"""
vllm/v1/engine/core.py
+8
-1
@@ -114,7 +114,14 @@ class EngineCore:
114
num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(
115
vllm_config
116
)
117
-
117
+ if kv_cache_config.kv_cache_groups:
118
+ vllm_config.cache_config.block_size = min(
119
+ g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups
120
+ )
121
+ elif vllm_config.cache_config.block_size is None:
122
+ # Attention-free models (encoder-only, SSM) — use default.
123
+ vllm_config.cache_config.block_size = 16
124
+ vllm_config.validate_block_size()
125
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
126
vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
127
self.collective_rpc("initialize_cache", args=(num_gpu_blocks, num_cpu_blocks))
vllm/v1/executor/multiproc_executor.py
+4
@@ -41,6 +41,7 @@ from vllm.distributed.parallel_state import (
41
)
42
from vllm.envs import enable_envs_cache
43
from vllm.logger import init_logger
44
+from vllm.platforms import current_platform
45
from vllm.tracing import instrument, maybe_init_worker_tracer
46
from vllm.utils.network_utils import (
47
get_distributed_init_method,
@@ -579,6 +580,9 @@ class WorkerProc:
580
self._init_message_queues(input_shm_handle, vllm_config)
581
self.worker.load_model()
582
583
+ # Set block size based on the attention backends
584
+ current_platform.update_block_size_for_backend(vllm_config)
585
+
586
# Enable environment variable cache (e.g. assume no more
587
# environment variable overrides after this point)
588
enable_envs_cache()
vllm/v1/executor/ray_executor.py
+5
@@ -385,6 +385,11 @@ class RayDistributedExecutor(Executor):
385
self.collective_rpc("init_device")
386
self.collective_rpc("load_model")
387
388
+ def _update_block_size(worker):
389
+ current_platform.update_block_size_for_backend(worker.vllm_config)
390
+
391
+ self.collective_rpc(_update_block_size)
392
+
393
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
394
self.pp_tp_workers.append([])
395
for tp_rank in range(self.parallel_config.tensor_parallel_size):
vllm/v1/executor/uniproc_executor.py
+2
@@ -12,6 +12,7 @@ import torch.distributed as dist
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
+from vllm.platforms import current_platform
16
from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
17
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
18
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
@@ -46,6 +47,7 @@ class UniProcExecutor(Executor):
47
self.driver_worker.init_worker(all_kwargs=[kwargs])
48
self.driver_worker.init_device()
49
self.driver_worker.load_model()
50
+ current_platform.update_block_size_for_backend(self.vllm_config)
51
52
def _distributed_args(self) -> tuple[str, int, int]:
53
"""Return (distributed_init_method, rank, local_rank)."""
vllm/v1/worker/gpu_model_runner.py
+3
-2
@@ -513,6 +513,7 @@ class GPUModelRunner(
513
custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (
514
tuple(logits_processors) if logits_processors is not None else ()
515
)
516
+ placeholder_block_size = self.cache_config.block_size or 16
517
self.input_batch = InputBatch(
518
max_num_reqs=self.max_num_reqs,
519
# We need to use the encoder length for encoder-decoer
@@ -522,8 +523,8 @@ class GPUModelRunner(
523
device=self.device,
524
pin_memory=self.pin_memory,
525
vocab_size=self.model_config.get_vocab_size(),
525
- block_sizes=[self.cache_config.block_size],
526
- kernel_block_sizes=[self.cache_config.block_size],
526
+ block_sizes=[placeholder_block_size],
527
+ kernel_block_sizes=[placeholder_block_size],
528
is_spec_decode=bool(self.vllm_config.speculative_config),
529
logitsprocs=build_logitsprocs(
530
self.vllm_config,