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

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

Reapply [Attention] Refactor `check_and_update_config` (#35122)

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

Matthew Bonanni committed Mar 9, 2026 at 10:17 UTC 77a73458e3ae8b5b7a2a13f78d3a6b4d39b1414d

32 files changed +311 -279

tests/kernels/attention/test_attention_selector.py

+40 -35

@@ -6,7 +6,12 @@ from unittest.mock import patch

6 import pytest

7 import torch

9 -from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config

9 +from vllm.config import (

10 + AttentionConfig,

11 + CacheConfig,

12 + VllmConfig,

13 + set_current_vllm_config,

14 +)

15 from vllm.platforms import current_platform

16 from vllm.platforms.cpu import CpuPlatform

17 from vllm.platforms.cuda import CudaPlatform

@@ -84,12 +89,15 @@ def test_backend_selection(

89 """Test attention backend selection with valid device-backend pairs."""

90 # Create AttentionConfig with the specified backend

91 attention_config = AttentionConfig(backend=AttentionBackendEnum[name])

87 - vllm_config = VllmConfig(attention_config=attention_config)

92 + cache_config = CacheConfig(block_size=block_size)

93 + vllm_config = VllmConfig(

94 + attention_config=attention_config, cache_config=cache_config

95 + )

97 with set_current_vllm_config(vllm_config):

98 if device == "cpu":

99 with patch("vllm.platforms.current_platform", CpuPlatform()):

92 - backend = get_attn_backend(16, torch.float16, None, block_size)

100 + backend = get_attn_backend(16, torch.float16, None)

101 assert backend.get_name() == "CPU_ATTN"

102

103 elif device == "hip":

@@ -104,20 +112,16 @@ def test_backend_selection(

112 if name == "TRITON_MLA" and block_size == 1:

113 # TRITON_MLA doesn't support block_size == 1

114 with pytest.raises(ValueError):

107 - get_attn_backend(

108 - 576, torch.float16, None, block_size, use_mla=use_mla

109 - )

115 + get_attn_backend(576, torch.float16, None, use_mla=use_mla)

116 else:

117 # Valid backend-block_size combination

118 backend = get_attn_backend(

113 - 576, torch.float16, None, block_size, use_mla=use_mla

119 + 576, torch.float16, None, use_mla=use_mla

120 )

121 expected = name

122 assert backend.get_name() == expected

123 else:

118 - backend = get_attn_backend(

119 - 32, torch.float16, None, block_size, use_mla=use_mla

120 - )

124 + backend = get_attn_backend(32, torch.float16, None, use_mla=use_mla)

125 expected = "ROCM_ATTN"

126 assert backend.get_name() == expected

127

@@ -141,7 +145,7 @@ def test_backend_selection(

145 if capability[0] != 10:

146 pytest.skip("CUTLASS MLA is not supported on this platform")

147 backend = get_attn_backend(

144 - 576, torch.float16, None, block_size, use_mla=use_mla

148 + 576, torch.float16, None, use_mla=use_mla

149 )

150 expected = "CUTLASS_MLA"

151 assert backend.get_name() == expected

@@ -156,7 +160,7 @@ def test_backend_selection(

160 "FlashInfer MLA only supports block_size 32 or 64"

161 )

162 backend = get_attn_backend(

159 - 576, torch.float16, None, block_size, use_mla=use_mla

163 + 576, torch.float16, None, use_mla=use_mla

164 )

165 expected = "FLASHINFER_MLA"

166 assert backend.get_name() == expected

@@ -175,7 +179,6 @@ def test_backend_selection(

179 576,

180 torch.float16,

181 None,

178 - block_size,

182 use_mla=use_mla,

183 )

184 expected = name

@@ -190,27 +193,23 @@ def test_backend_selection(

193 "FlashAttention MLA not supported on this platform"

194 )

195 backend = get_attn_backend(

193 - 576, torch.float16, None, block_size, use_mla=use_mla

196 + 576, torch.float16, None, use_mla=use_mla

197 )

198 expected = "FLASH_ATTN_MLA"

199 assert backend.get_name() == expected

200 else:

201 # TRITON_MLA or other fallback

202 backend = get_attn_backend(

200 - 576, torch.float16, None, block_size, use_mla=use_mla

203 + 576, torch.float16, None, use_mla=use_mla

204 )

205 expected = "TRITON_MLA"

206 assert backend.get_name() == expected

207 elif name == "FLASHINFER":

205 - backend = get_attn_backend(

206 - 64, torch.float16, None, block_size, use_mla=use_mla

207 - )

208 + backend = get_attn_backend(64, torch.float16, None, use_mla=use_mla)

209 expected = "FLASHINFER"

210 assert backend.get_name() == expected

211 elif name == "FLASH_ATTN":

211 - backend = get_attn_backend(

212 - 32, torch.float16, None, block_size, use_mla=use_mla

213 - )

212 + backend = get_attn_backend(32, torch.float16, None, use_mla=use_mla)

213 expected = "FLASH_ATTN"

214 assert backend.get_name() == expected

215

@@ -224,12 +223,12 @@ def test_fp32_fallback(device: str):

223 with set_current_vllm_config(vllm_config):

224 if device == "cpu":

225 with patch("vllm.platforms.current_platform", CpuPlatform()):

227 - backend = get_attn_backend(16, torch.float32, None, 16)

226 + backend = get_attn_backend(16, torch.float32, None)

227 assert backend.get_name() == "CPU_ATTN"

228

229 elif device == "cuda":

230 with patch("vllm.platforms.current_platform", CudaPlatform()):

232 - backend = get_attn_backend(16, torch.float32, None, 16)

231 + backend = get_attn_backend(16, torch.float32, None)

232 assert backend.get_name() == "FLEX_ATTENTION"

233

234

@@ -241,35 +240,40 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):

240 )

241

242 attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN)

244 - vllm_config = VllmConfig(attention_config=attention_config)

243 + cache_config = CacheConfig(block_size=16)

244 + vllm_config = VllmConfig(

245 + attention_config=attention_config, cache_config=cache_config

246 + )

247

248 with set_current_vllm_config(vllm_config):

249 # Unsupported CUDA arch

250 monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))

249 - backend = get_attn_backend(16, torch.float16, None, 16)

251 + backend = get_attn_backend(16, torch.float16, None)

252 assert backend.get_name() != "FLASH_ATTN"

253

254 # Reset the monkeypatch for subsequent tests

255 monkeypatch.undo()

256

257 # Unsupported data type

256 - backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16)

258 + backend = get_attn_backend(16, torch.float8_e4m3fn, None)

259 assert backend.get_name() != "FLASH_ATTN"

260

261 # Unsupported kv cache data type

260 - backend = get_attn_backend(16, torch.float16, "fp8", 16)

262 + backend = get_attn_backend(16, torch.float16, "fp8")

263 assert backend.get_name() != "FLASH_ATTN"

264

265 # Unsupported block size

264 - backend = get_attn_backend(16, torch.float16, None, 8)

266 + vllm_config.cache_config.block_size = 8

267 + backend = get_attn_backend(16, torch.float16, None)

268 assert backend.get_name() != "FLASH_ATTN"

269

270 # flash-attn is not installed

271 import sys

272

273 + vllm_config.cache_config.block_size = 16

274 original_module = sys.modules.get("vllm_flash_attn")

275 monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)

272 - backend = get_attn_backend(16, torch.float16, None, 16)

276 + backend = get_attn_backend(16, torch.float16, None)

277 assert backend.get_name() != "FLASH_ATTN"

278

279 # Restore the original module if it existed

@@ -279,7 +283,7 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):

283 monkeypatch.delitem(sys.modules, "vllm_flash_attn", raising=False)

284

285 # Unsupported head size

282 - backend = get_attn_backend(17, torch.float16, None, 16)

286 + backend = get_attn_backend(17, torch.float16, None)

287 assert backend.get_name() != "FLASH_ATTN"

288

289

@@ -320,7 +324,7 @@ def test_auto_backend_selection_behavior():

324 set_current_vllm_config(vllm_config_auto),

325 patch("vllm.platforms.current_platform", CpuPlatform()),

326 ):

323 - backend_auto = get_attn_backend(16, torch.float16, None, 16)

327 + backend_auto = get_attn_backend(16, torch.float16, None)

328

329 _cached_get_attn_backend.cache_clear()

330

@@ -328,7 +332,7 @@ def test_auto_backend_selection_behavior():

332 set_current_vllm_config(vllm_config_none),

333 patch("vllm.platforms.current_platform", CpuPlatform()),

334 ):

331 - backend_none = get_attn_backend(16, torch.float16, None, 16)

335 + backend_none = get_attn_backend(16, torch.float16, None)

336

337 # Both should select the same backend

338 assert backend_auto.get_name() == backend_none.get_name()

@@ -358,7 +362,10 @@ def test_per_head_quant_scales_backend_selection(

362 backend=AttentionBackendEnum[backend_name],

363 flash_attn_version=flash_attn_version,

364 )

361 - vllm_config = VllmConfig(attention_config=attention_config)

365 + cache_config = CacheConfig(block_size=64)

366 + vllm_config = VllmConfig(

367 + attention_config=attention_config, cache_config=cache_config

368 + )

369

370 with (

371 set_current_vllm_config(vllm_config),

@@ -376,7 +383,6 @@ def test_per_head_quant_scales_backend_selection(

383 head_size=128,

384 dtype=torch.float16,

385 kv_cache_dtype="fp8",

379 - block_size=64,

386 use_per_head_quant_scales=True,

387 )

388 assert backend.get_name() == backend_name

@@ -386,7 +392,6 @@ def test_per_head_quant_scales_backend_selection(

392 head_size=128,

393 dtype=torch.float16,

394 kv_cache_dtype="fp8",

389 - block_size=64,

395 use_per_head_quant_scales=True,

396 )

397 assert backend_name in str(exc_info.value)

tests/models/multimodal/processing/test_tensor_schema.py

+4 -1

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

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/v1/spec_decode/test_eagle.py

+1 -1

@@ -80,7 +80,7 @@ def _create_proposer(

80 device = current_platform.device_type

81 vllm_config = VllmConfig(

82 model_config=model_config,

83 - cache_config=CacheConfig(),

83 + cache_config=CacheConfig(block_size=16),

84 speculative_config=speculative_config,

85 device_config=DeviceConfig(device=device),

86 parallel_config=ParallelConfig(),

vllm/config/cache.py

+26 -8

@@ -2,16 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from dataclasses import field

5 -from typing import Literal

5 +from typing import ClassVar, Literal

7 -from pydantic import Field, SkipValidation, field_validator

7 +from pydantic import Field, SkipValidation, field_validator, model_validator

9 from vllm.config.utils import config

10 from vllm.logger import init_logger

12 logger = init_logger(__name__)

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

14 CacheDType = Literal[

15 "auto",

16 "bfloat16",

@@ -31,12 +30,13 @@ KVOffloadingBackend = Literal["native", "lmcache"]

30 class CacheConfig:

31 """Configuration for the KV cache."""

34 - block_size: SkipValidation[BlockSize] = None # type: ignore[assignment]

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

33 + DEFAULT_BLOCK_SIZE: ClassVar[int] = 16

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

38 - be set in `Platform.check_and_update_config()` based on the current

39 - platform."""

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

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

37 + Accepts None (meaning "use default"). After construction, always int."""

38 + user_specified_block_size: bool = field(default=False, init=False)

39 + """Whether block_size was explicitly provided. Derived automatically."""

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

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

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

@@ -169,6 +169,8 @@ class CacheConfig:

169 "prefix_caching_hash_algo",

170 "cpu_kvcache_space_bytes",

171 "mamba_page_size_padded",

172 + "user_specified_block_size",

173 + "_block_size_resolved",

174 # Post-init/derived counters

175 "num_gpu_blocks",

176 "num_cpu_blocks",

@@ -186,6 +188,22 @@ class CacheConfig:

188 # metrics info

189 return {key: str(value) for key, value in self.__dict__.items()}

190

191 + _block_size_resolved: bool = field(default=False, init=False)

192 + """Guard against pydantic re-running _apply_block_size_default."""

193 +

194 + @model_validator(mode="after")

195 + def _apply_block_size_default(self) -> "CacheConfig":

196 + # Pydantic re-runs validators when CacheConfig is nested inside

197 + # another pydantic model (e.g. VllmConfig). Guard against that.

198 + if self._block_size_resolved:

199 + return self

200 + object.__setattr__(self, "_block_size_resolved", True)

201 + if self.block_size is None:

202 + object.__setattr__(self, "block_size", self.DEFAULT_BLOCK_SIZE)

203 + else:

204 + object.__setattr__(self, "user_specified_block_size", True)

205 + return self

206 +

207 @field_validator("cache_dtype", mode="after")

208 @classmethod

209 def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:

vllm/config/vllm.py

+47 -46

@@ -1026,32 +1026,6 @@ class VllmConfig:

1026 )

1027 current_platform.check_and_update_config(self)

1028

1029 - # If DCP, ensure the block size is right.

1030 - if self.parallel_config.decode_context_parallel_size > 1:

1031 - if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (

1032 - self.parallel_config.cp_kv_cache_interleave_size

1033 - != self.parallel_config.dcp_kv_cache_interleave_size

1034 - ):

1035 - self.parallel_config.cp_kv_cache_interleave_size = (

1036 - self.parallel_config.dcp_kv_cache_interleave_size

1037 - )

1038 - logger.warning_once(

1039 - "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"

1040 - "_interleave_size. And dcp-kv-cache-interleave-size will be "

1041 - "deprecated when PCP is fully supported."

1042 - )

1043 - assert (

1044 - self.parallel_config.cp_kv_cache_interleave_size

1045 - <= self.cache_config.block_size

1046 - and self.cache_config.block_size

1047 - % self.parallel_config.cp_kv_cache_interleave_size

1048 - == 0

1049 - ), (

1050 - f"Block_size({self.cache_config.block_size}) should be greater "

1051 - "than or equal to and divisible by cp_kv_cache_interleave_size "

1052 - f"({self.parallel_config.cp_kv_cache_interleave_size})."

1053 - )

1054 -

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

1030 effective_dp_size = (

1031 self.parallel_config.data_parallel_size

@@ -1219,26 +1193,6 @@ class VllmConfig:

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

1194 self.scheduler_config.disable_hybrid_kv_cache_manager = False

1195

1222 - if self.cache_config.mamba_cache_mode == "align":

1223 - assert (

1224 - self.cache_config.block_size

1225 - <= self.scheduler_config.max_num_batched_tokens

1226 - ), (

1227 - "In Mamba cache align mode, block_size "

1228 - f"({self.cache_config.block_size}) must be <= "

1229 - "max_num_batched_tokens "

1230 - f"({self.scheduler_config.max_num_batched_tokens})."

1231 - )

1232 - if self.scheduler_config.long_prefill_token_threshold > 0:

1233 - assert (

1234 - self.scheduler_config.long_prefill_token_threshold

1235 - >= self.cache_config.block_size

1236 - )

1237 - assert not self.scheduler_config.disable_chunked_mm_input, (

1238 - "Chunked MM input is required because we need the flexibility to "

1239 - "schedule a multiple of block_size tokens even if they are in the "

1240 - "middle of a mm input"

1241 - )

1196 if self.compilation_config.debug_dump_path:

1197 self.compilation_config.debug_dump_path = (

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

@@ -1673,6 +1627,53 @@ class VllmConfig:

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

1628 )

1629

1630 + def validate_block_size(self) -> None:

1631 + """Validate block_size against DCP and mamba constraints.

1632 +

1633 + Called after Platform.update_block_size_for_backend() has

1634 + finalised block_size.

1635 + """

1636 + block_size = self.cache_config.block_size

1637 +

1638 + # DCP interleave-size compatibility

1639 + if self.parallel_config.decode_context_parallel_size > 1:

1640 + if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (

1641 + self.parallel_config.cp_kv_cache_interleave_size

1642 + != self.parallel_config.dcp_kv_cache_interleave_size

1643 + ):

1644 + self.parallel_config.cp_kv_cache_interleave_size = (

1645 + self.parallel_config.dcp_kv_cache_interleave_size

1646 + )

1647 + logger.warning_once(

1648 + "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"

1649 + "_interleave_size. And dcp-kv-cache-interleave-size will be "

1650 + "deprecated when PCP is fully supported."

1651 + )

1652 + assert (

1653 + self.parallel_config.cp_kv_cache_interleave_size <= block_size

1654 + and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0

1655 + ), (

1656 + f"Block_size({block_size}) should be greater "

1657 + "than or equal to and divisible by cp_kv_cache_interleave_size "

1658 + f"({self.parallel_config.cp_kv_cache_interleave_size})."

1659 + )

1660 +

1661 + # Mamba cache align-mode constraints

1662 + if self.cache_config.mamba_cache_mode == "align":

1663 + assert block_size <= self.scheduler_config.max_num_batched_tokens, (

1664 + "In Mamba cache align mode, block_size "

1665 + f"({block_size}) must be <= "

1666 + "max_num_batched_tokens "

1667 + f"({self.scheduler_config.max_num_batched_tokens})."

1668 + )

1669 + if self.scheduler_config.long_prefill_token_threshold > 0:

1670 + assert self.scheduler_config.long_prefill_token_threshold >= block_size

1671 + assert not self.scheduler_config.disable_chunked_mm_input, (

1672 + "Chunked MM input is required because we need the flexibility "

1673 + "to schedule a multiple of block_size tokens even if they are "

1674 + "in the middle of a mm input"

1675 + )

1676 +

1677 @model_validator(mode="after")

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

1679 if self.model_config is None:

vllm/distributed/kv_transfer/kv_connector/utils.py

-1

@@ -500,7 +500,6 @@ def get_current_attn_backend(vllm_config: VllmConfig):

500 head_size=vllm_config.model_config.get_head_size(),

501 dtype=vllm_config.model_config.dtype,

502 kv_cache_dtype=vllm_config.cache_config.cache_dtype,

503 - block_size=vllm_config.cache_config.block_size,

503 use_mla=vllm_config.model_config.use_mla,

504 )

505 return backend

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

-1

@@ -726,7 +726,6 @@ class MoRIIOConnectorWorker:

726 self.model_config.get_head_size(),

727 self.model_config.dtype,

728 self.cache_config.cache_dtype,

729 - self.block_size,

729 use_mla=self.use_mla,

730 )

731

vllm/engine/arg_utils.py

+2 -3

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

62 get_attr_docs,

63 )

64 from vllm.config.cache import (

65 - BlockSize,

65 CacheDType,

66 KVOffloadingBackend,

67 MambaCacheMode,

@@ -440,7 +439,7 @@ class EngineArgs:

439 max_parallel_loading_workers: int | None = (

440 ParallelConfig.max_parallel_loading_workers

441 )

443 - block_size: BlockSize = CacheConfig.block_size

442 + block_size: int | None = None

443 enable_prefix_caching: bool | None = None

444 prefix_caching_hash_algo: PrefixCachingHashAlgo = (

445 CacheConfig.prefix_caching_hash_algo

@@ -1521,7 +1520,7 @@ class EngineArgs:

1520 )

1521

1522 cache_config = CacheConfig(

1524 - block_size=self.block_size,

1523 + block_size=self.block_size, # type: ignore[arg-type]

1524 gpu_memory_utilization=self.gpu_memory_utilization,

1525 kv_cache_memory_bytes=self.kv_cache_memory_bytes,

1526 cache_dtype=resolved_cache_dtype, # type: ignore[arg-type]

vllm/model_executor/layers/attention/attention.py

-3

@@ -221,11 +221,9 @@ class Attention(nn.Module, AttentionLayerBase):

221 vllm_config = get_current_vllm_config()

222 if cache_config is not None:

223 kv_cache_dtype = cache_config.cache_dtype

224 - block_size = cache_config.block_size

224 calculate_kv_scales = cache_config.calculate_kv_scales

225 else:

226 kv_cache_dtype = "auto"

228 - block_size = 16

227 calculate_kv_scales = False

228

229 # llm-compressor mdls need to set cache_dtype to "fp8" manually.

@@ -275,7 +273,6 @@ class Attention(nn.Module, AttentionLayerBase):

273 head_size,

274 dtype,

275 kv_cache_dtype,

278 - block_size,

276 use_mla=False,

277 has_sink=self.has_sink,

278 use_mm_prefix=self.use_mm_prefix,

vllm/model_executor/layers/attention/chunked_local_attention.py

+6 -9

@@ -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}_"

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

@@ -94,16 +95,12 @@ class ChunkedLocalAttention(Attention):

95 dtype = torch.get_default_dtype()

96 if cache_config is not None:

97 kv_cache_dtype = cache_config.cache_dtype

97 - block_size = cache_config.block_size

98 else:

99 kv_cache_dtype = "auto"

100 - block_size = 16

100

102 - underlying_attn_backend = get_attn_backend(

103 - head_size, dtype, kv_cache_dtype, block_size

104 - )

101 + underlying_attn_backend = get_attn_backend(head_size, dtype, kv_cache_dtype)

102 attn_backend = create_chunked_local_attention_backend(

106 - underlying_attn_backend, attention_chunk_size, block_size

103 + underlying_attn_backend, attention_chunk_size

104 )

105

106 super().__init__(

vllm/model_executor/layers/attention/cross_attention.py

-3

@@ -188,10 +188,8 @@ class CrossAttention(Attention):

188

189 if cache_config is not None:

190 kv_cache_dtype = cache_config.cache_dtype

191 - block_size = cache_config.block_size

191 else:

192 kv_cache_dtype = "auto"

194 - block_size = 16

193

194 if attn_type is not None:

195 assert attn_type == AttentionType.ENCODER_DECODER, (

@@ -202,7 +200,6 @@ class CrossAttention(Attention):

200 head_size,

201 dtype,

202 kv_cache_dtype,

205 - block_size,

203 attn_type=AttentionType.ENCODER_DECODER,

204 )

205 attn_backend = create_cross_attention_backend(underlying_attn_backend)

vllm/model_executor/layers/attention/encoder_only_attention.py

-3

@@ -66,16 +66,13 @@ class EncoderOnlyAttention(Attention):

67 if cache_config is not None:

68 kv_cache_dtype = cache_config.cache_dtype

69 - block_size = cache_config.block_size

69 else:

70 kv_cache_dtype = "auto"

72 - block_size = 16

72 underlying_attn_backend = get_attn_backend(

73 head_size,

74 dtype,

75 kv_cache_dtype,

78 - block_size,

76 attn_type=AttentionType.ENCODER_ONLY,

77 )

vllm/model_executor/layers/attention/mla_attention.py

+12 -8

@@ -323,11 +323,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):

323

324 if cache_config is not None:

325 kv_cache_dtype = cache_config.cache_dtype

326 - block_size = cache_config.block_size

326 calculate_kv_scales = cache_config.calculate_kv_scales

327 else:

328 kv_cache_dtype = "auto"

330 - block_size = 16

329 calculate_kv_scales = False

330 self.quant_config = quant_config

331

@@ -336,7 +334,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

334 self.head_size,

335 dtype,

336 kv_cache_dtype,

339 - block_size,

337 use_mla=True,

338 use_sparse=use_sparse,

339 num_heads=self.num_heads,

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

446 )

447

448 # Attributes for forward_impl method

452 - self.chunked_prefill_workspace_size = (

453 - MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(

454 - get_current_vllm_config()

455 - )

456 - )

449 + self._vllm_config = get_current_vllm_config()

450 + self._chunked_prefill_workspace_size: int | None = None

451 self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

452 static=True,

453 group_shape=GroupShape.PER_TENSOR,

454 compile_native=True,

455 )

456

457 + @property

458 + def chunked_prefill_workspace_size(self) -> int:

459 + if self._chunked_prefill_workspace_size is None:

460 + self._chunked_prefill_workspace_size = (

461 + MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(

462 + self._vllm_config

463 + )

464 + )

465 + return self._chunked_prefill_workspace_size

466 +

467 def forward(

468 self,

469 q: torch.Tensor,

vllm/model_executor/layers/attention/static_sink_attention.py

+3 -8

@@ -126,17 +126,13 @@ class StaticSinkAttention(Attention, CustomOp):

126

127 if cache_config is not None:

128 kv_cache_dtype = cache_config.cache_dtype

129 - block_size = cache_config.block_size

129 else:

130 kv_cache_dtype = "auto"

132 - block_size = 16

131

132 if attn_backend is not None:

133 underlying_attn_backend = attn_backend

134 else:

137 - underlying_attn_backend = get_attn_backend(

138 - head_size, dtype, kv_cache_dtype, block_size

139 - )

135 + underlying_attn_backend = get_attn_backend(head_size, dtype, kv_cache_dtype)

136 attn_backend = create_static_sink_attention_backend(

137 underlying_attn_backend, # type: ignore[arg-type]

138 sink_len=sink_len,

@@ -153,7 +149,6 @@ class StaticSinkAttention(Attention, CustomOp):

149 CustomOp.__init__(self)

150

151 self.sink_len = sink_len

156 - self.block_size = block_size

152 self.sink_populated = False

153 self.sink_key = None

154 self.sink_value = None

@@ -212,12 +207,12 @@ class StaticSinkAttention(Attention, CustomOp):

207

208 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

209 # Block size may get updated after model loading, refresh it

215 - block_size = vllm_config.cache_config.block_size

210 + self.block_size = vllm_config.cache_config.block_size

211 # Should not be called for enc-dec or encoder-only attention.

212 assert self.attn_type == AttentionType.DECODER

213

214 return SinkFullAttentionSpec(

220 - block_size=block_size,

215 + block_size=self.block_size,

216 num_kv_heads=self.num_kv_heads,

217 head_size=self.head_size,

218 head_size_v=self.head_size_v,

vllm/model_executor/models/config.py

+3 -4

@@ -217,10 +217,9 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):

217 mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token

218 )

219

220 - # override attention block size if either (a) the

221 - # user has not set it or (b) the user has set it

222 - # too small.

223 - if cache_config.block_size is None or cache_config.block_size < attn_block_size:

220 + # override attention block size if it is too small,

221 + # even if the user has explicitly set it

222 + if cache_config.block_size < attn_block_size:

223 cache_config.block_size = attn_block_size

224 logger.info(

225 "Setting attention block size to %d tokens "

vllm/model_executor/models/whisper_causal.py

-3

@@ -290,16 +290,13 @@ class WhisperCausalAttentionWithBlockPooling(Attention):

290

291 if cache_config is not None:

292 kv_cache_dtype = cache_config.cache_dtype

293 - block_size = cache_config.block_size

293 else:

294 kv_cache_dtype = "auto"

296 - block_size = 16

295

296 underlying_attn_backend = get_attn_backend(

297 head_size,

298 dtype,

299 kv_cache_dtype,

302 - block_size,

300 attn_type=attn_type,

301 )

302 attn_backend = create_whisper_attention_backend_with_block_pooling(

vllm/platforms/cpu.py

+7 -1

@@ -185,7 +185,7 @@ class CpuPlatform(Platform):

185

186 cache_config = vllm_config.cache_config

187

188 - if cache_config.block_size is None:

188 + if not cache_config.user_specified_block_size:

189 cache_config.block_size = 128

190

191 if cache_config.block_size % 32 != 0:

@@ -361,6 +361,12 @@ class CpuPlatform(Platform):

361 vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,

362 )

363

364 + @classmethod

365 + def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

366 + # TODO: CPU still sets block_size in check_and_update_config.

367 + # Move that logic here so block_size is chosen by the backend.

368 + pass

369 +

370 @classmethod

371 def get_allowed_cpu_core_node_list(cls) -> tuple[list[int], list[LogicalCPUInfo]]:

372 assert platform.system() == "Linux"

vllm/platforms/cuda.py

+29 -117

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

166

167 @classmethod

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

169 - from vllm.v1.attention.backends.registry import AttentionBackendEnum

170 -

169 parallel_config = vllm_config.parallel_config

170 model_config = vllm_config.model_config

171

172 if parallel_config.worker_cls == "auto":

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

174

177 - cache_config = vllm_config.cache_config

178 - if cache_config and cache_config.block_size is None:

179 - cache_config.block_size = 16

180 -

181 - # TODO(lucas): handle this more gracefully

182 - # Note: model_config may be None during testing

183 - # Note: block_size is initialized in

184 - # HybridAttentionMambaModelConfig.verify_and_update_config

185 - # for models with both attention and mamba,

186 - # and doesn't need to be reinitialized here

187 - if (

188 - model_config is not None

189 - and model_config.use_mla

190 - and cache_config.block_size is not None

191 - ):

192 - use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")

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

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

195 - # else we default to CutlassMLA. For each case, we force the

196 - # required block_size.

197 - use_flashmla = False

198 - use_cutlass_mla = False

199 - use_flashinfer_mla = False

200 - use_flashmla_sparse = False

201 - use_flashinfer_mla_sparse = False

202 -

203 - from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

204 -

205 - if vllm_config.attention_config.backend is None:

206 - # Default case

207 - hf_text_config = model_config.hf_text_config

208 - qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)

209 - if (

210 - cls.is_device_capability_family(100)

211 - and not use_sparse

212 - and qk_nope_head_dim == 128

213 - ):

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

215 - # and only if qk_nope_head_dim == 128 (kernel constraint)

216 - use_flashinfer_mla = True

217 - # Set the backend in AttentionConfig so it's used during

218 - # backend selection

219 - vllm_config.attention_config.backend = (

220 - AttentionBackendEnum.FLASHINFER_MLA

221 - )

222 - elif cls.is_device_capability_family(100) and not use_sparse:

223 - # Fall back to CUTLASS_MLA as 2nd priority on Blackwell

224 - use_cutlass_mla = True

225 - elif is_flashmla_dense_supported()[0]:

226 - # Non-Blackwell with FlashMLA support

227 - use_flashmla = True

228 - else:

229 - # Fallback: will use Triton MLA or other compatible backend

230 - pass

231 - else:

232 - # Forced case

233 - backend = vllm_config.attention_config.backend

234 - use_flashmla = backend == AttentionBackendEnum.FLASHMLA

235 - use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA

236 - use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA

237 - use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE

238 - use_flashinfer_mla_sparse = (

239 - backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE

240 - )

241 -

242 - if (

243 - use_flashmla

244 - and is_flashmla_dense_supported()[0]

245 - and cache_config.block_size % 64 != 0

246 - ):

247 - cache_config.block_size = 64

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

249 -

250 - if use_cutlass_mla and cache_config.block_size % 128 != 0:

251 - cache_config.block_size = 128

252 - logger.info(

253 - "Forcing kv cache block size to 128 for CUTLASS_MLA backend."

254 - )

255 -

256 - if (

257 - use_flashinfer_mla

258 - and cache_config.block_size != 32

259 - and cache_config.block_size % 64 != 0

260 - ):

261 - cache_config.block_size = 64

262 - logger.info(

263 - "Forcing kv cache block size to 64 for FlashInferMLA backend."

264 - )

265 -

266 - if use_sparse:

267 - if not (use_flashmla_sparse or use_flashinfer_mla_sparse):

268 - use_flashmla_sparse = True

269 -

270 - if use_flashmla_sparse and cache_config.block_size != 64:

271 - cache_config.block_size = 64

272 - logger.info(

273 - "Forcing kv cache block size to 64 for FlashMLASparse backend."

274 - )

275 - elif use_flashinfer_mla_sparse and cache_config.block_size not in (

276 - 32,

277 - 64,

278 - ):

279 - cache_config.block_size = 64

280 - logger.info(

281 - "Forcing kv cache block size to 64 for FlashInferMLASparse "

282 - "backend."

283 - )

284 -

175 scheduler_config = vllm_config.scheduler_config

176 # Note: model_config may be None during testing

177 if (

@@ -312,10 +202,10 @@ class CudaPlatformBase(Platform):

202 num_heads: int | None = None,

203 ) -> tuple[

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

315 - dict["AttentionBackendEnum", list[str]],

205 + dict["AttentionBackendEnum", tuple[int, list[str]]],

206 ]:

207 valid_backends_priorities = []

318 - invalid_reasons = {}

208 + invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}

209

210 backend_priorities = _get_backend_priorities(

211 attn_selector_config.use_mla,

@@ -332,7 +222,7 @@ class CudaPlatformBase(Platform):

222 except ImportError:

223 invalid_reasons_i = ["ImportError"]

224 if invalid_reasons_i:

335 - invalid_reasons[backend] = invalid_reasons_i

225 + invalid_reasons[backend] = (priority, invalid_reasons_i)

226 else:

227 valid_backends_priorities.append((backend, priority))

228

@@ -341,14 +231,13 @@ class CudaPlatformBase(Platform):

231 @classmethod

232 def get_attn_backend_cls(

233 cls,

344 - selected_backend: "AttentionBackendEnum",

234 + selected_backend: "AttentionBackendEnum | None",

235 attn_selector_config: "AttentionSelectorConfig",

236 num_heads: int | None = None,

237 ) -> str:

238 device_capability = cls.get_device_capability()

239 assert device_capability is not None

240

351 - attn_selector_config = attn_selector_config._replace(block_size=None)

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

242 if selected_backend is not None:

243 try:

@@ -370,7 +259,7 @@ class CudaPlatformBase(Platform):

259

260 # No selected backend or the selected backend is invalid,

261 # so we try finding a valid backend.

373 - valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

262 + valid_backends_priorities, all_invalid_reasons = cls.get_valid_backends(

263 device_capability=device_capability,

264 attn_selector_config=attn_selector_config,

265 num_heads=num_heads,

@@ -379,7 +268,7 @@ class CudaPlatformBase(Platform):

268 "{"

269 + ", ".join(

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

382 - for backend, reasons in invalid_reasons.items()

271 + for backend, (_, reasons) in all_invalid_reasons.items()

272 )

273 + "}"

274 )

@@ -402,6 +291,29 @@ class CudaPlatformBase(Platform):

291 )

292 selected_index = sorted_indices[0]

293 selected_backend = valid_backends_priorities[selected_index][0]

294 + selected_priority = valid_backends_priorities[selected_index][1]

295 +

296 + # If the user specified --block-size (but not --attention-backend),

297 + # check whether that constraint precluded any higher-priority backends.

298 + if attn_selector_config.block_size is not None:

299 + excluded = [

300 + backend

301 + for backend, (priority, reasons) in all_invalid_reasons.items()

302 + if priority < selected_priority

303 + and reasons == ["block_size not supported"]

304 + ]

305 + if excluded:

306 + names = ", ".join(b.name for b in excluded)

307 + logger.warning(

308 + "--block-size %d precluded higher-priority backend(s) "

309 + "%s. Using %s instead, which may result in reduced "

310 + "performance. Consider removing --block-size to "

311 + "auto-select the optimal block size.",

312 + attn_selector_config.block_size,

313 + names,

314 + selected_backend.name,

315 + )

316 +

317 logger.info_once(

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

319 selected_backend.name,

vllm/platforms/interface.py

+50

@@ -420,6 +420,56 @@ class Platform:

420 """

421 pass

422

423 + @classmethod

424 + def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

425 + """

426 + Ensure block_size is compatible with the attention backend.

427 + """

428 + from vllm.config.cache import CacheConfig

429 +

430 + cache_config = vllm_config.cache_config

431 + if cache_config.user_specified_block_size:

432 + # User specified --block-size; keep it.

433 + return

434 +

435 + model_config = vllm_config.model_config

436 + # model_config may be None during testing.

437 + # Skip hybrid models — their block_size is managed by

438 + # HybridAttentionMambaModelConfig.

439 + if model_config is None or model_config.is_hybrid:

440 + cache_config.block_size = CacheConfig.DEFAULT_BLOCK_SIZE

441 + return

442 +

443 + from vllm.config.vllm import (

444 + get_layers_from_vllm_config,

445 + set_current_vllm_config,

446 + )

447 + from vllm.model_executor.layers.attention_layer_base import (

448 + AttentionLayerBase,

449 + )

450 +

451 + attn_layers = get_layers_from_vllm_config(

452 + vllm_config,

453 + AttentionLayerBase, # type: ignore[type-abstract]

454 + )

455 + if not attn_layers:

456 + cache_config.block_size = CacheConfig.DEFAULT_BLOCK_SIZE

457 + return

458 +

459 + first_layer = next(iter(attn_layers.values()))

460 + backend_cls = first_layer.get_attn_backend()

461 + with set_current_vllm_config(vllm_config):

462 + preferred = backend_cls.get_preferred_block_size(

463 + CacheConfig.DEFAULT_BLOCK_SIZE

464 + )

465 + if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:

466 + logger.info(

467 + "Setting kv cache block size to %d for %s backend.",

468 + preferred,

469 + backend_cls.get_name(),

470 + )

471 + cache_config.block_size = preferred

472 +

473 @classmethod

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

475 """

vllm/platforms/rocm.py

+7 -1

@@ -687,7 +687,7 @@ class RocmPlatform(Platform):

687 )

688 compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE

689

690 - if cache_config and cache_config.block_size is None:

690 + if cache_config and not cache_config.user_specified_block_size:

691 if (

692 envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION and envs.VLLM_ROCM_USE_AITER

693 # NOTE: This block has been deprecated

@@ -707,6 +707,12 @@ class RocmPlatform(Platform):

707 if parallel_config.worker_cls == "auto":

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

709

710 + @classmethod

711 + def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

712 + # TODO: ROCm still sets block_size in check_and_update_config.

713 + # Move that logic here so block_size is chosen by the backend.

714 + pass

715 +

716 @classmethod

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

718 if model_arch in _ROCM_UNSUPPORTED_MODELS:

vllm/platforms/xpu.py

+7 -1

@@ -162,7 +162,7 @@ class XPUPlatform(Platform):

162 model_config = vllm_config.model_config

163 parallel_config = vllm_config.parallel_config

164 # in V1(or with chunked prefill) block_size is 64

165 - if cache_config and cache_config.block_size is None:

165 + if cache_config and not cache_config.user_specified_block_size:

166 cache_config.block_size = 64

167

168 # lazy import to avoid circular import

@@ -227,6 +227,12 @@ class XPUPlatform(Platform):

227 # ref. https://openucx.readthedocs.io/en/master/faq.html

228 os.environ["UCX_MEMTYPE_CACHE"] = "n"

229

230 + @classmethod

231 + def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

232 + # TODO: XPU still sets block_size in check_and_update_config.

233 + # Move that logic here so block_size is chosen by the backend.

234 + pass

235 +

236 @classmethod

237 def support_hybrid_kv_cache(cls) -> bool:

238 return True

vllm/v1/attention/backend.py

+14 -9

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

@@ -210,7 +215,7 @@ class AttentionBackend(ABC):

215 head_size: int,

216 dtype: torch.dtype,

217 kv_cache_dtype: "CacheDType | None",

213 - block_size: int,

218 + block_size: int | None,

219 use_mla: bool,

220 has_sink: bool,

221 use_sparse: bool,

@@ -224,7 +229,7 @@ class AttentionBackend(ABC):

229 head_size: int,

230 dtype: torch.dtype,

231 kv_cache_dtype: "CacheDType | None",

227 - block_size: int,

232 + block_size: int | None,

233 use_mla: bool,

234 has_sink: bool,

235 use_sparse: bool,

vllm/v1/attention/backends/mla/flashattn_mla.py

+1 -1

@@ -75,7 +75,7 @@ class FlashAttnMLABackend(MLACommonBackend):

75 head_size: int,

76 dtype: torch.dtype,

77 kv_cache_dtype: CacheDType | None,

78 - block_size: int,

78 + block_size: int | None,

79 use_mla: bool,

80 has_sink: bool,

81 use_sparse: bool,

vllm/v1/attention/backends/mla/flashinfer_mla.py

+1 -1

@@ -69,7 +69,7 @@ class FlashInferMLABackend(MLACommonBackend):

69 head_size: int,

70 dtype: torch.dtype,

71 kv_cache_dtype: CacheDType | None,

72 - block_size: int,

72 + block_size: int | None,

73 use_mla: bool,

74 has_sink: bool,

75 use_sparse: bool,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

+1 -1

@@ -106,7 +106,7 @@ class FlashInferMLASparseBackend(AttentionBackend):

106 head_size: int,

107 dtype: torch.dtype,

108 kv_cache_dtype: CacheDType | None,

109 - block_size: int,

109 + block_size: int | None,

110 use_mla: bool,

111 has_sink: bool,

112 use_sparse: bool,

vllm/v1/attention/backends/mla/flashmla.py

+1 -1

@@ -80,7 +80,7 @@ class FlashMLABackend(MLACommonBackend):

80 head_size: int,

81 dtype: torch.dtype,

82 kv_cache_dtype: CacheDType | None,

83 - block_size: int,

83 + block_size: int | None,

84 use_mla: bool,

85 has_sink: bool,

86 use_sparse: bool,

vllm/v1/attention/selector.py

+6 -1

@@ -49,7 +49,6 @@ def get_attn_backend(

49 head_size: int,

50 dtype: torch.dtype,

51 kv_cache_dtype: str | None,

52 - block_size: int | None,

52 use_mla: bool = False,

53 has_sink: bool = False,

54 use_sparse: bool = False,

@@ -71,6 +70,12 @@ def get_attn_backend(

71 vllm_config = get_current_vllm_config()

73 + cache_config = vllm_config.cache_config

74 + if cache_config is not None and cache_config.user_specified_block_size:

75 + block_size = cache_config.block_size

76 + else:

77 + block_size = None

78 +

79 attn_selector_config = AttentionSelectorConfig(

80 head_size=head_size,

81 dtype=dtype,

vllm/v1/engine/core.py

+5 -1

@@ -122,7 +122,11 @@ class EngineCore:

122 num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(

123 vllm_config

124 )

125 -

125 + if kv_cache_config.kv_cache_groups:

126 + vllm_config.cache_config.block_size = min(

127 + g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups

128 + )

129 + vllm_config.validate_block_size()

130 vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks

131 vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks

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

vllm/v1/executor/multiproc_executor.py

@@ -42,6 +42,7 @@ from vllm.distributed.parallel_state import (

42 )

43 from vllm.envs import enable_envs_cache

44 from vllm.logger import init_logger

45 +from vllm.platforms import current_platform

46 from vllm.tracing import instrument, maybe_init_worker_tracer

47 from vllm.utils.network_utils import (

48 get_distributed_init_method,

@@ -617,6 +618,9 @@ class WorkerProc:

618 )

619 self.worker.load_model()

620

621 + # Set block size based on the attention backends

622 + current_platform.update_block_size_for_backend(vllm_config)

623 +

624 # Initialize message queues after init_device() since multi-node setups

625 # (nnodes_within_dp > 1) require distributed groups to be initialized

626 self._init_message_queues(input_shm_handle, vllm_config)

vllm/v1/executor/ray_executor.py

@@ -387,6 +387,11 @@ class RayDistributedExecutor(Executor):

387 self.collective_rpc("init_device")

388 self.collective_rpc("load_model")

389

390 + def _update_block_size(worker):

391 + current_platform.update_block_size_for_backend(worker.vllm_config)

392 +

393 + self.collective_rpc(_update_block_size)

394 +

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

396 self.pp_tp_workers.append([])

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

vllm/v1/executor/uniproc_executor.py

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

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.executor.abstract import Executor

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

48 if not is_eep_new_worker:

49 self.driver_worker.init_device()

50 self.driver_worker.load_model()

51 + current_platform.update_block_size_for_backend(self.vllm_config)

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

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

vllm/v1/worker/gpu_model_runner.py

+27 -7

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

32 set_current_vllm_config,

33 update_config,

34 )

35 +from vllm.config.cache import CacheConfig

36 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer

37 from vllm.distributed.eplb.eplb_state import EplbState

38 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group

@@ -586,6 +587,11 @@ class GPUModelRunner(

587 custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (

588 tuple(logits_processors) if logits_processors is not None else ()

589 )

590 + placeholder_block_size = (

591 + self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE

592 + )

593 + self._init_block_sizes = [placeholder_block_size]

594 + self._init_kernel_block_sizes = [placeholder_block_size]

595 self.input_batch = InputBatch(

596 max_num_reqs=self.max_num_reqs,

597 # We need to use the encoder length for encoder-decoder

@@ -595,8 +601,8 @@ class GPUModelRunner(

601 device=self.device,

602 pin_memory=self.pin_memory,

603 vocab_size=self.model_config.get_vocab_size(),

598 - block_sizes=[self.cache_config.block_size],

599 - kernel_block_sizes=[self.cache_config.block_size],

604 + block_sizes=[placeholder_block_size],

605 + kernel_block_sizes=[placeholder_block_size],

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

607 logitsprocs=build_logitsprocs(

608 self.vllm_config,

@@ -6112,8 +6118,10 @@ class GPUModelRunner(

6118 ) -> None:

6119 """

6120 Re-initialize the input batch if the block sizes are different from

6115 - `[self.cache_config.block_size]`. This usually happens when there

6116 - are multiple KV cache groups.

6121 + what it was originally created with. This happens when the final

6122 + block size (determined after model loading) differs from the

6123 + placeholder used during __init__, or when there are multiple

6124 + KV cache groups.

6125

6126 Args:

6127 kv_cache_config: The KV cache configuration.

@@ -6138,14 +6146,17 @@ class GPUModelRunner(

6146 ) + kv_cache_group.kv_cache_spec.num_speculative_blocks

6147 max_num_blocks.append(max_num_blocks_per_req)

6148

6141 - if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [

6142 - self.cache_config.block_size

6143 - ]:

6149 + if (

6150 + block_sizes != self._init_block_sizes

6151 + or kernel_block_sizes != self._init_kernel_block_sizes

6152 + ):

6153 assert self.offload_config.uva.cpu_offload_gb == 0, (

6154 "Cannot re-initialize the input batch when CPU weight "

6155 "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501

6156 "for more details."

6157 )

6158 + self._init_block_sizes = block_sizes

6159 + self._init_kernel_block_sizes = kernel_block_sizes

6160 self.input_batch = InputBatch(

6161 max_num_reqs=self.max_num_reqs,

6162 max_model_len=max_model_len,

@@ -6162,6 +6173,15 @@ class GPUModelRunner(

6173 is_pooling_model=self.is_pooling_model,

6174 )

6175

6176 + assert self._init_block_sizes == block_sizes, (

6177 + f"InputBatch block_sizes {self._init_block_sizes} != "

6178 + f"kv_cache block_sizes {block_sizes}"

6179 + )

6180 + assert self._init_kernel_block_sizes == kernel_block_sizes, (

6181 + f"InputBatch kernel_block_sizes {self._init_kernel_block_sizes} "

6182 + f"!= kv_cache kernel_block_sizes {kernel_block_sizes}"

6183 + )

6184 +

6185 def _allocate_kv_cache_tensors(

6186 self, kv_cache_config: KVCacheConfig

6187 ) -> dict[str, torch.Tensor]: