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

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, 03:45 UTC

[Attention] Refactor `check_and_update_config` (#33600)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Feb 17, 2026 at 20:06 UTC 7743152957236f21fc36f0402f9678159976ccc5

4 files changed +270 -172

vllm/config/cache.py

+4 -7

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

20 logger = init_logger(__name__)

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

22 CacheDType = Literal[

23 "auto",

24 "bfloat16",

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

38 class CacheConfig:

39 """Configuration for the KV cache."""

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.

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

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

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."""

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."""

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/engine/arg_utils.py

+1 -2

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

59 get_attr_docs,

60 )

61 from vllm.config.cache import (

62 - BlockSize,

62 CacheDType,

63 KVOffloadingBackend,

64 MambaCacheMode,

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

430 max_parallel_loading_workers: int | None = (

431 ParallelConfig.max_parallel_loading_workers

432 )

434 - block_size: BlockSize = CacheConfig.block_size

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

434 enable_prefix_caching: bool | None = None

435 prefix_caching_hash_algo: PrefixCachingHashAlgo = (

436 CacheConfig.prefix_caching_hash_algo

vllm/platforms/cuda.py

+253 -156

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

166 parallel_config = vllm_config.parallel_config

167 model_config = vllm_config.model_config

168

@@ -172,112 +170,19 @@ class CudaPlatformBase(Platform):

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

171

172 cache_config = vllm_config.cache_config

175 - if cache_config and cache_config.block_size is None:

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

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

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

187 scheduler_config = vllm_config.scheduler_config

188 # Note: model_config may be None during testing

@@ -293,6 +198,150 @@ class CudaPlatformBase(Platform):

198 )

199 scheduler_config.disable_chunked_mm_input = True

200

201 + @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

218 + cache_config = vllm_config.cache_config

219 +

220 + device_capability = cls.get_device_capability()

221 + if device_capability is None:

222 + return

223 +

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,

234 + )

235 +

236 + user_specified_backend = vllm_config.attention_config.backend

237 + num_heads = model_config.get_num_attention_heads(

238 + vllm_config.parallel_config,

239 + )

240 + 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,

248 + )

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

344 +

345 @classmethod

346 def get_current_memory_usage(

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

@@ -336,77 +385,125 @@ class CudaPlatformBase(Platform):

385 return valid_backends_priorities, invalid_reasons

386

387 @classmethod

339 - def get_attn_backend_cls(

388 + def select_attention_backend(

389 cls,

341 - selected_backend: "AttentionBackendEnum",

390 + selected_backend: "AttentionBackendEnum | None",

391 attn_selector_config: "AttentionSelectorConfig",

392 + device_capability: "DeviceCapability",

393 + raise_on_invalid: bool = True,

394 num_heads: int | None = None,

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)

395 + ) -> "AttentionBackendEnum | None":

396 + """Select the best attention backend for the given configuration.

397 +

398 + Args:

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

400 + attn_selector_config: Configuration for attention selection

401 + device_capability: Device capability info

402 + raise_on_invalid: If True, raise ValueError when no valid backend

403 + num_heads: Number of attention heads per GPU, used for backend

404 + priority ordering on Blackwell GPUs

405 +

406 + Returns:

407 + The selected backend enum, or None if no valid backend found

408 + and raise_on_invalid is False

409 + """

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

411 if selected_backend is not None:

412 try:

413 backend_class = selected_backend.get_class()

353 - invalid_reasons = backend_class.validate_configuration(

414 + validation_errors = backend_class.validate_configuration(

415 device_capability=device_capability,

416 **attn_selector_config._asdict(),

417 )

418 except ImportError:

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

419 + validation_errors = ["ImportError"]

420 + if validation_errors:

421 + if raise_on_invalid:

422 + raise ValueError(

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

424 + f"this configuration. Reason: {validation_errors}"

425 + )

426 + return None

427 + return selected_backend

428

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

369 - # so we try finding a valid backend.

429 + # No selected backend, so find the best valid one.

430 valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

431 device_capability=device_capability,

432 attn_selector_config=attn_selector_config,

433 num_heads=num_heads,

434 )

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

435 +

436 if len(valid_backends_priorities) == 0:

389 - raise ValueError(

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

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

392 - )

437 + if raise_on_invalid:

438 + reasons_str = (

439 + "{"

440 + + ", ".join(

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

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

443 + )

444 + + "}"

445 + )

446 + config_str = attn_selector_config.__repr__()

447 + raise ValueError(

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

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

450 + )

451 + return None

452

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",

453 + # Select the one with the highest priority (lowest index).

454 + sorted_backends = sorted(valid_backends_priorities, key=lambda x: x[1])

455 + return sorted_backends[0][0]

456 +

457 + @classmethod

458 + def get_attn_backend_cls(

459 + cls,

460 + selected_backend: "AttentionBackendEnum | None",

461 + attn_selector_config: "AttentionSelectorConfig",

462 + num_heads: int | None = None,

463 + ) -> str:

464 + device_capability = cls.get_device_capability()

465 + assert device_capability is not None

466 +

467 + chosen_backend = cls.select_attention_backend(

468 + selected_backend=selected_backend,

469 + attn_selector_config=attn_selector_config,

470 + num_heads=num_heads,

471 + device_capability=device_capability,

472 + raise_on_invalid=True,

473 )

474 + assert chosen_backend is not None # raise_on_invalid=True guarantees this

475 +

476 + # Log the selection

477 + if selected_backend is not None:

478 + logger.info("Using %s backend.", chosen_backend)

479 + else:

480 + # Get all valid backends for logging

481 + valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

482 + device_capability=device_capability,

483 + attn_selector_config=attn_selector_config,

484 + num_heads=num_heads,

485 + )

486 + reasons_str = (

487 + "{"

488 + + ", ".join(

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

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

491 + )

492 + + "}"

493 + )

494 + config_str = attn_selector_config.__repr__()

495 + logger.debug_once(

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

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

498 + )

499 + logger.info_once(

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

501 + chosen_backend.name,

502 + tuple(b[0].name for b in valid_backends_priorities),

503 + scope="local",

504 + )

505

409 - return selected_backend.get_path()

506 + return chosen_backend.get_path()

507

508 @classmethod

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

vllm/v1/attention/backend.py

+12 -7

@@ -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, get_args

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

9 import numpy as np

10 import torch

@@ -144,15 +144,9 @@ 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 -

147 if block_size is None:

148 return True

149

152 - valid_sizes = get_args(BlockSize)

153 - if block_size not in valid_sizes:

154 - return False

155 -

150 supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()

151 if not supported_kernel_block_sizes:

152 return True

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

161 return True

162 return False

163

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 +

175 @classmethod

176 def is_mla(cls) -> bool:

177 return False