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

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 8, 00:00 UTC

[CI] Revert PRs 34818 and 33600 (#34979)

Lucas Wilkinson committed Feb 20, 2026 at 16:25 UTC aaefc58ee0f023ec7bd3671ca83aae1b8a8f271d

16 files changed +242 -294

tests/models/multimodal/processing/test_tensor_schema.py

+1 -4

@@ -13,7 +13,6 @@ import torch.nn as nn

13 from PIL import Image

15 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config

16 -from vllm.config.cache import CacheConfig

16 from vllm.config.multimodal import (

17 AudioDummyOptions,

18 BaseDummyOptions,

@@ -132,9 +131,7 @@ def initialize_dummy_model(

131 ):

132 temp_file = tempfile.mkstemp()[1]

133 current_device = torch.get_default_device()

135 - vllm_config = VllmConfig(

136 - model_config=model_config, cache_config=CacheConfig(block_size=16)

137 - )

134 + vllm_config = VllmConfig(model_config=model_config)

135 with set_current_vllm_config(vllm_config=vllm_config):

136 init_distributed_environment(

137 world_size=1,

tests/models/utils.py

+2 -7

@@ -457,9 +457,6 @@ 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)

460 num_experts = n_group * 2 if n_group is not None else 2

461

462 # we use three layers for Gemma-3n to check

@@ -489,14 +486,12 @@ def dummy_hf_overrides(

486 # Only set MoE related config when the model has MoE layers.

487 # Otherwise all models detected as MoE by _get_transformers_backend_cls.

488 if model_arch_config.num_experts > 0:

492 - orig_topk = getattr(text_config, "num_experts_per_tok", 2)

493 - topk = min(orig_topk, 2)

489 update_dict.update(

490 {

491 "num_experts": num_experts,

497 - "num_experts_per_tok": topk,

492 + "num_experts_per_tok": 2,

493 # Kimi uses `num_experts_per_token`.

499 - "num_experts_per_token": topk,

494 + "num_experts_per_token": 2,

495 "num_local_experts": num_experts,

496 # Otherwise there will not be any expert layers

497 "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(block_size=16),

81 + cache_config=CacheConfig(),

82 speculative_config=speculative_config,

83 device_config=DeviceConfig(device=device),

84 parallel_config=ParallelConfig(),

vllm/config/cache.py

+7 -4

@@ -19,6 +19,7 @@ else:

20 logger = init_logger(__name__)

22 +BlockSize = Literal[1, 8, 16, 32, 64, 128, 256]

23 CacheDType = Literal[

24 "auto",

25 "bfloat16",

@@ -38,11 +39,13 @@ KVOffloadingBackend = Literal["native", "lmcache"]

39 class CacheConfig:

40 """Configuration for the KV cache."""

41 - block_size: SkipValidation[int] = None # type: ignore[assignment]

42 - """Size of a contiguous cache block in number of tokens.

42 + block_size: SkipValidation[BlockSize] = None # type: ignore[assignment]

43 + """Size of a contiguous cache block in number of tokens. On CUDA devices,

44 + only block sizes up to 32 are supported.

44 - This is None until the platform sets it. Always an int by the time

45 - the engine starts."""

46 + This config has no static default. If left unspecified by the user, it will

47 + be set in `Platform.check_and_update_config()` based on the current

48 + platform."""

49 gpu_memory_utilization: float = Field(default=0.9, gt=0, le=1)

50 """The fraction of GPU memory to be used for the model executor, which can

51 range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory

vllm/config/vllm.py

+46 -51

@@ -915,6 +915,32 @@ 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 +

944 # Do this after all the updates to compilation_config.mode

945 effective_dp_size = (

946 self.parallel_config.data_parallel_size

@@ -1082,6 +1108,26 @@ class VllmConfig:

1108 # Default to enable HMA if not explicitly disabled by user or logic above.

1109 self.scheduler_config.disable_hybrid_kv_cache_manager = False

1110

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 + )

1131 if self.compilation_config.debug_dump_path:

1132 self.compilation_config.debug_dump_path = (

1133 self.compilation_config.debug_dump_path.absolute().expanduser()

@@ -1442,57 +1488,6 @@ class VllmConfig:

1488 f"compilation_config={self.compilation_config!r}"

1489 )

1490

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 -

1491 @model_validator(mode="after")

1492 def validate_mamba_block_size(self) -> "VllmConfig":

1493 if self.model_config is None:

vllm/engine/arg_utils.py

+2 -1

@@ -59,6 +59,7 @@ from vllm.config import (

59 get_attr_docs,

60 )

61 from vllm.config.cache import (

62 + BlockSize,

63 CacheDType,

64 KVOffloadingBackend,

65 MambaCacheMode,

@@ -430,7 +431,7 @@ class EngineArgs:

431 max_parallel_loading_workers: int | None = (

432 ParallelConfig.max_parallel_loading_workers

433 )

433 - block_size: int = None # type: ignore[assignment]

434 + block_size: BlockSize = CacheConfig.block_size

435 enable_prefix_caching: bool | None = None

436 prefix_caching_hash_algo: PrefixCachingHashAlgo = (

437 CacheConfig.prefix_caching_hash_algo

vllm/model_executor/layers/attention/chunked_local_attention.py

+5 -6

@@ -30,8 +30,9 @@ 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,

34 ) -> type[AttentionBackend]:

34 - prefix = f"ChunkedLocalAttention_{attention_chunk_size}_"

35 + prefix = f"ChunkedLocalAttention_{attention_chunk_size}_{block_size}_"

37 underlying_builder = underlying_attn_backend.get_builder_cls()

38 assert issubclass(underlying_builder, AttentionMetadataBuilder)

@@ -54,9 +55,7 @@ def create_chunked_local_attention_backend(

55 fast_build: bool = False,

56 ):

57 cm, make_virtual_batches_block_table = make_local_attention_virtual_batches(

57 - attention_chunk_size,

58 - common_attn_metadata,

59 - self.kv_cache_spec.block_size,

58 + attention_chunk_size, common_attn_metadata, block_size

59 )

60 metadata = super().build(common_prefix_len, cm, fast_build)

61 metadata.make_virtual_batches_block_table = make_virtual_batches_block_table

@@ -98,13 +97,13 @@ class ChunkedLocalAttention(Attention):

97 block_size = cache_config.block_size

98 else:

99 kv_cache_dtype = "auto"

101 - block_size = None

100 + block_size = 16

101

102 underlying_attn_backend = get_attn_backend(

103 head_size, dtype, kv_cache_dtype, block_size

104 )

105 attn_backend = create_chunked_local_attention_backend(

107 - underlying_attn_backend, attention_chunk_size

106 + underlying_attn_backend, attention_chunk_size, block_size

107 )

108

109 super().__init__(

vllm/model_executor/layers/attention/mla_attention.py

+5 -12

@@ -407,24 +407,17 @@ class MLAAttention(nn.Module, AttentionLayerBase):

407 )

408

409 # Attributes for forward_impl method

410 - self._vllm_config = get_current_vllm_config()

411 - self._chunked_prefill_workspace_size: int | None = None

410 + self.chunked_prefill_workspace_size = (

411 + MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(

412 + get_current_vllm_config()

413 + )

414 + )

415 self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

416 static=True,

417 group_shape=GroupShape.PER_TENSOR,

418 compile_native=True,

419 )

420

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 -

421 def forward(

422 self,

423 q: torch.Tensor,

vllm/platforms/cuda.py

+163 -167

@@ -163,12 +163,122 @@ class CudaPlatformBase(Platform):

163

164 @classmethod

165 def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:

166 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

167 +

168 parallel_config = vllm_config.parallel_config

169 model_config = vllm_config.model_config

170

171 if parallel_config.worker_cls == "auto":

172 parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"

173

174 + cache_config = vllm_config.cache_config

175 + if cache_config and cache_config.block_size is None:

176 + cache_config.block_size = 16

177 +

178 + # TODO(lucas): handle this more gracefully

179 + # Note: model_config may be None during testing

180 + # Note: block_size is initialized in

181 + # HybridAttentionMambaModelConfig.verify_and_update_config

182 + # for models with both attention and mamba,

183 + # and doesn't need to be reinitialized here

184 + if (

185 + model_config is not None

186 + and model_config.use_mla

187 + and cache_config.block_size is not None

188 + ):

189 + use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")

190 + # If `--attention-config.backend` is not set and we are using MLA,

191 + # then we default to FlashMLA backend for non-blackwell GPUs,

192 + # else we default to CutlassMLA. For each case, we force the

193 + # required block_size.

194 + use_flashmla = False

195 + use_cutlass_mla = False

196 + use_flashinfer_mla = False

197 + use_flashmla_sparse = False

198 + use_flashinfer_mla_sparse = False

199 +

200 + from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

201 +

202 + if vllm_config.attention_config.backend is None:

203 + # Default case

204 + hf_text_config = model_config.hf_text_config

205 + qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)

206 + if (

207 + cls.is_device_capability_family(100)

208 + and not use_sparse

209 + and qk_nope_head_dim == 128

210 + ):

211 + # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2)

212 + # and only if qk_nope_head_dim == 128 (kernel constraint)

213 + use_flashinfer_mla = True

214 + # Set the backend in AttentionConfig so it's used during

215 + # backend selection

216 + vllm_config.attention_config.backend = (

217 + AttentionBackendEnum.FLASHINFER_MLA

218 + )

219 + elif cls.is_device_capability_family(100) and not use_sparse:

220 + # Fall back to CUTLASS_MLA as 2nd priority on Blackwell

221 + use_cutlass_mla = True

222 + elif is_flashmla_dense_supported()[0]:

223 + # Non-Blackwell with FlashMLA support

224 + use_flashmla = True

225 + else:

226 + # Fallback: will use Triton MLA or other compatible backend

227 + pass

228 + else:

229 + # Forced case

230 + backend = vllm_config.attention_config.backend

231 + use_flashmla = backend == AttentionBackendEnum.FLASHMLA

232 + use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA

233 + use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA

234 + use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE

235 + use_flashinfer_mla_sparse = (

236 + backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE

237 + )

238 +

239 + if (

240 + use_flashmla

241 + and is_flashmla_dense_supported()[0]

242 + and cache_config.block_size % 64 != 0

243 + ):

244 + cache_config.block_size = 64

245 + logger.info("Forcing kv cache block size to 64 for FlashMLA backend.")

246 +

247 + if use_cutlass_mla and cache_config.block_size % 128 != 0:

248 + cache_config.block_size = 128

249 + logger.info(

250 + "Forcing kv cache block size to 128 for CUTLASS_MLA backend."

251 + )

252 +

253 + if (

254 + use_flashinfer_mla

255 + and cache_config.block_size != 32

256 + and cache_config.block_size % 64 != 0

257 + ):

258 + cache_config.block_size = 64

259 + logger.info(

260 + "Forcing kv cache block size to 64 for FlashInferMLA backend."

261 + )

262 +

263 + if use_sparse:

264 + if not (use_flashmla_sparse or use_flashinfer_mla_sparse):

265 + use_flashmla_sparse = True

266 +

267 + if use_flashmla_sparse and cache_config.block_size != 64:

268 + cache_config.block_size = 64

269 + logger.info(

270 + "Forcing kv cache block size to 64 for FlashMLASparse backend."

271 + )

272 + elif use_flashinfer_mla_sparse and cache_config.block_size not in (

273 + 32,

274 + 64,

275 + ):

276 + cache_config.block_size = 64

277 + logger.info(

278 + "Forcing kv cache block size to 64 for FlashInferMLASparse "

279 + "backend."

280 + )

281 +

282 scheduler_config = vllm_config.scheduler_config

283 # Note: model_config may be None during testing

284 if (

@@ -183,49 +293,6 @@ class CudaPlatformBase(Platform):

293 )

294 scheduler_config.disable_chunked_mm_input = True

295

186 - @classmethod

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 -

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 -

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 -

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):

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 - )

227 - cache_config.block_size = preferred

228 -

296 @classmethod

297 def get_current_memory_usage(

298 cls, device: torch.types.Device | None = None

@@ -242,10 +309,10 @@ class CudaPlatformBase(Platform):

309 num_heads: int | None = None,

310 ) -> tuple[

311 list[tuple["AttentionBackendEnum", int]],

245 - dict["AttentionBackendEnum", tuple[int, list[str]]],

312 + dict["AttentionBackendEnum", list[str]],

313 ]:

314 valid_backends_priorities = []

248 - invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}

315 + invalid_reasons = {}

316

317 backend_priorities = _get_backend_priorities(

318 attn_selector_config.use_mla,

@@ -262,155 +329,84 @@ class CudaPlatformBase(Platform):

329 except ImportError:

330 invalid_reasons_i = ["ImportError"]

331 if invalid_reasons_i:

265 - invalid_reasons[backend] = (priority, invalid_reasons_i)

332 + invalid_reasons[backend] = invalid_reasons_i

333 else:

334 valid_backends_priorities.append((backend, priority))

335

336 return valid_backends_priorities, invalid_reasons

337

338 @classmethod

272 - def select_attention_backend(

339 + def get_attn_backend_cls(

340 cls,

274 - selected_backend: "AttentionBackendEnum | None",

341 + selected_backend: "AttentionBackendEnum",

342 attn_selector_config: "AttentionSelectorConfig",

276 - device_capability: "DeviceCapability",

277 - raise_on_invalid: bool = True,

343 num_heads: int | None = None,

279 - ) -> "AttentionBackendEnum | None":

280 - """Select the best attention backend for the given configuration.

281 -

282 - Args:

283 - selected_backend: User-specified backend, or None for auto-selection

284 - attn_selector_config: Configuration for attention selection

285 - device_capability: Device capability info

286 - raise_on_invalid: If True, raise ValueError when no valid backend

287 - num_heads: Number of attention heads per GPU, used for backend

288 - priority ordering on Blackwell GPUs

289 -

290 - Returns:

291 - The selected backend enum, or None if no valid backend found

292 - and raise_on_invalid is False

293 - """

344 + ) -> str:

345 + device_capability = cls.get_device_capability()

346 + assert device_capability is not None

347 +

348 + attn_selector_config = attn_selector_config._replace(block_size=None)

349 # First try checking just the selected backend, if there is one.

350 if selected_backend is not None:

351 try:

352 backend_class = selected_backend.get_class()

298 - validation_errors = backend_class.validate_configuration(

353 + invalid_reasons = backend_class.validate_configuration(

354 device_capability=device_capability,

355 **attn_selector_config._asdict(),

356 )

357 except ImportError:

303 - validation_errors = ["ImportError"]

304 - if validation_errors:

305 - if raise_on_invalid:

306 - raise ValueError(

307 - f"Selected backend {selected_backend} is not valid for "

308 - f"this configuration. Reason: {validation_errors}"

309 - )

310 - return None

311 - return selected_backend

358 + invalid_reasons = ["ImportError"]

359 + if invalid_reasons:

360 + raise ValueError(

361 + f"Selected backend {selected_backend} is not valid for "

362 + f"this configuration. Reason: {invalid_reasons}"

363 + )

364 + else:

365 + logger.info("Using %s backend.", selected_backend)

366 + return selected_backend.get_path()

367

313 - # No selected backend, so find the best valid one.

368 + # No selected backend or the selected backend is invalid,

369 + # so we try finding a valid backend.

370 valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

371 device_capability=device_capability,

372 attn_selector_config=attn_selector_config,

373 num_heads=num_heads,

374 )

319 -

375 + reasons_str = (

376 + "{"

377 + + ", ".join(

378 + f"{backend.name}: [{', '.join(reasons)}]"

379 + for backend, reasons in invalid_reasons.items()

380 + )

381 + + "}"

382 + )

383 + config_str = attn_selector_config.__repr__()

384 + logger.debug_once(

385 + f"Some attention backends are not valid for {cls.device_name} with "

386 + f"{config_str}. Reasons: {reasons_str}."

387 + )

388 if len(valid_backends_priorities) == 0:

321 - if raise_on_invalid:

322 - reasons_str = (

323 - "{"

324 - + ", ".join(

325 - f"{backend.name}: [{', '.join(reasons)}]"

326 - for backend, (_, reasons) in invalid_reasons.items()

327 - )

328 - + "}"

329 - )

330 - config_str = attn_selector_config.__repr__()

331 - raise ValueError(

332 - f"No valid attention backend found for {cls.device_name} "

333 - f"with {config_str}. Reasons: {reasons_str}."

334 - )

335 - return None

336 -

337 - # Select the one with the highest priority (lowest index).

338 - sorted_backends = sorted(valid_backends_priorities, key=lambda x: x[1])

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(

366 - cls,

367 - selected_backend: "AttentionBackendEnum | None",

368 - attn_selector_config: "AttentionSelectorConfig",

369 - num_heads: int | None = None,

370 - ) -> str:

371 - device_capability = cls.get_device_capability()

372 - assert device_capability is not None

389 + raise ValueError(

390 + f"No valid attention backend found for {cls.device_name} "

391 + f"with {config_str}. Reasons: {reasons_str}."

392 + )

393

374 - chosen_backend = cls.select_attention_backend(

375 - selected_backend=selected_backend,

376 - attn_selector_config=attn_selector_config,

377 - num_heads=num_heads,

378 - device_capability=device_capability,

379 - raise_on_invalid=True,

394 + # We have found some valid backends. Select the one with the

395 + # highest priority.

396 + sorted_indices = sorted(

397 + range(len(valid_backends_priorities)),

398 + key=lambda i: valid_backends_priorities[i][1],

399 + )

400 + selected_index = sorted_indices[0]

401 + selected_backend = valid_backends_priorities[selected_index][0]

402 + logger.info_once(

403 + "Using %s attention backend out of potential backends: %s.",

404 + selected_backend.name,

405 + "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",

406 + scope="local",

407 )

381 - assert chosen_backend is not None # raise_on_invalid=True guarantees this

382 -

383 - # Log the selection

384 - if selected_backend is not None:

385 - logger.info("Using %s backend.", chosen_backend)

386 - else:

387 - # Get all valid backends for logging

388 - valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

389 - device_capability=device_capability,

390 - attn_selector_config=attn_selector_config,

391 - num_heads=num_heads,

392 - )

393 - reasons_str = (

394 - "{"

395 - + ", ".join(

396 - f"{backend.name}: [{', '.join(reasons)}]"

397 - for backend, (_, reasons) in invalid_reasons.items()

398 - )

399 - + "}"

400 - )

401 - config_str = attn_selector_config.__repr__()

402 - logger.debug_once(

403 - f"Some attention backends are not valid for {cls.device_name} with "

404 - f"{config_str}. Reasons: {reasons_str}."

405 - )

406 - logger.info_once(

407 - "Using %s attention backend out of potential backends: %s",

408 - chosen_backend.name,

409 - tuple(backend.name for backend, _ in valid_backends_priorities),

410 - scope="local",

411 - )

408

413 - return chosen_backend.get_path()

409 + return selected_backend.get_path()

410

411 @classmethod

412 def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:

vllm/platforms/interface.py

-7

@@ -406,13 +406,6 @@ 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 -

409 @classmethod

410 def verify_model_arch(cls, model_arch: str) -> None:

411 """

vllm/v1/attention/backend.py

+7 -12

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

4 from abc import ABC, abstractmethod

5 from dataclasses import dataclass, replace

6 from enum import Enum

7 -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar

7 +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar, get_args

9 import numpy as np

10 import torch

@@ -144,9 +144,15 @@ class AttentionBackend(ABC):

144

145 @classmethod

146 def supports_block_size(cls, block_size: int | None) -> bool:

147 + from vllm.config.cache import BlockSize

148 +

149 if block_size is None:

150 return True

151

152 + valid_sizes = get_args(BlockSize)

153 + if block_size not in valid_sizes:

154 + return False

155 +

156 supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()

157 if not supported_kernel_block_sizes:

158 return True

@@ -161,17 +167,6 @@ class AttentionBackend(ABC):

167 return True

168 return False

169

164 - @classmethod

165 - def get_preferred_block_size(cls, default_block_size: int = 16) -> int:

166 - supported_sizes = cls.get_supported_kernel_block_sizes()

167 - if not supported_sizes:

168 - return default_block_size

169 -

170 - if cls.supports_block_size(default_block_size):

171 - return default_block_size

172 -

173 - return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)

174 -

170 @classmethod

171 def is_mla(cls) -> bool:

172 return False

vllm/v1/engine/core.py

+1 -8

@@ -114,14 +114,7 @@ class EngineCore:

114 num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(

115 vllm_config

116 )

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()

117 +

118 vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks

119 vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks

120 self.collective_rpc("initialize_cache", args=(num_gpu_blocks, num_cpu_blocks))

vllm/v1/executor/multiproc_executor.py

-4

@@ -41,7 +41,6 @@ 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

44 from vllm.tracing import instrument, maybe_init_worker_tracer

45 from vllm.utils.network_utils import (

46 get_distributed_init_method,

@@ -580,9 +579,6 @@ class WorkerProc:

579 self._init_message_queues(input_shm_handle, vllm_config)

580 self.worker.load_model()

581

583 - # Set block size based on the attention backends

584 - current_platform.update_block_size_for_backend(vllm_config)

585 -

582 # Enable environment variable cache (e.g. assume no more

583 # environment variable overrides after this point)

584 enable_envs_cache()

vllm/v1/executor/ray_executor.py

-5

@@ -385,11 +385,6 @@ 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 -

388 for pp_rank in range(self.parallel_config.pipeline_parallel_size):

389 self.pp_tp_workers.append([])

390 for tp_rank in range(self.parallel_config.tensor_parallel_size):

vllm/v1/executor/uniproc_executor.py

-2

@@ -12,7 +12,6 @@ import torch.distributed as dist

13 import vllm.envs as envs

14 from vllm.logger import init_logger

15 -from vllm.platforms import current_platform

15 from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port

16 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

17 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

@@ -47,7 +46,6 @@ class UniProcExecutor(Executor):

46 self.driver_worker.init_worker(all_kwargs=[kwargs])

47 self.driver_worker.init_device()

48 self.driver_worker.load_model()

50 - current_platform.update_block_size_for_backend(self.vllm_config)

50 def _distributed_args(self) -> tuple[str, int, int]:

51 """Return (distributed_init_method, rank, local_rank)."""

vllm/v1/worker/gpu_model_runner.py

+2 -3

@@ -513,7 +513,6 @@ 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

516 self.input_batch = InputBatch(

517 max_num_reqs=self.max_num_reqs,

518 # We need to use the encoder length for encoder-decoer

@@ -523,8 +522,8 @@ class GPUModelRunner(

522 device=self.device,

523 pin_memory=self.pin_memory,

524 vocab_size=self.model_config.get_vocab_size(),

526 - block_sizes=[placeholder_block_size],

527 - kernel_block_sizes=[placeholder_block_size],

525 + block_sizes=[self.cache_config.block_size],

526 + kernel_block_sizes=[self.cache_config.block_size],

527 is_spec_decode=bool(self.vllm_config.speculative_config),

528 logitsprocs=build_logitsprocs(

529 self.vllm_config,