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

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

Improve configs - `CacheConfig` (#16835)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Apr 20, 2025 at 05:25 UTC 4b07d36891510ab43677695d2dbb098937eaa5e1

3 files changed +123 -157

vllm/config.py

+69 -51

@@ -1245,22 +1245,70 @@ class ModelConfig:

1245 or getattr(self.hf_config, "is_matryoshka", False))

1246

1247

1248 +BlockSize = Literal[8, 16, 32, 64, 128]

1249 +CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2"]

1250 +PrefixCachingHashAlgo = Literal["builtin", "sha256"]

1251 +

1252 +

1253 +@config

1254 +@dataclass

1255 class CacheConfig:

1249 - """Configuration for the KV cache.

1256 + """Configuration for the KV cache."""

1257

1251 - Args:

1252 - block_size: Size of a cache block in number of tokens.

1253 - gpu_memory_utilization: Fraction of GPU memory to use for the

1254 - vLLM execution.

1255 - swap_space: Size of the CPU swap space per GPU (in GiB).

1256 - cache_dtype: Data type for kv cache storage.

1257 - is_attention_free: Whether the model is attention-free.

1258 - num_gpu_blocks_override: Number of GPU blocks to use. This overrides the

1259 - profiled num_gpu_blocks if specified. Does nothing if None.

1260 - sliding_window: Sliding window size for the KV cache.

1261 - enable_prefix_caching: Whether to enable prefix caching.

1262 - cpu_offload_gb: Size of the CPU offload buffer in GiB.

1258 + block_size: Optional[BlockSize] = None

1259 + """Size of a contiguous cache block in number of tokens. This is ignored on

1260 + neuron devices and set to `--max-model-len`. On CUDA devices, only block

1261 + sizes up to 32 are supported. On HPU devices, block size defaults to 128.

1262 + """

1263 + gpu_memory_utilization: float = 0.9

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

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

1266 + utilization. If unspecified, will use the default value of 0.9. This is a

1267 + per-instance limit, and only applies to the current vLLM instance. It does

1268 + not matter if you have another vLLM instance running on the same GPU. For

1269 + example, if you have two vLLM instances running on the same GPU, you can

1270 + set the GPU memory utilization to 0.5 for each instance."""

1271 + swap_space: float = 4

1272 + """Size of the CPU swap space per GPU (in GiB)."""

1273 + cache_dtype: CacheDType = "auto"

1274 + """Data type for kv cache storage. If "auto", will use model data type.

1275 + CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports

1276 + fp8 (=fp8_e4m3)."""

1277 + is_attention_free: bool = False

1278 + """Whether the model is attention-free. This is primarily set in

1279 + `ModelConfig` and that value should be manually duplicated here."""

1280 + num_gpu_blocks_override: Optional[int] = None

1281 + """Number of GPU blocks to use. This overrides the profiled `num_gpu_blocks`

1282 + if specified. Does nothing if `None`. Used for testing preemption."""

1283 + sliding_window: Optional[int] = None

1284 + """Sliding window size for the KV cache. This is primarily set in

1285 + `ModelConfig` and that value should be manually duplicated here."""

1286 + enable_prefix_caching: Optional[bool] = None

1287 + """Whether to enable prefix caching. Disabled by default for V0. Enabled by

1288 + default for V1."""

1289 + prefix_caching_hash_algo: PrefixCachingHashAlgo = "builtin"

1290 + """Set the hash algorithm for prefix caching:\n

1291 + - "builtin" is Python's built-in hash.\n

1292 + - "sha256" is collision resistant but with certain overheads."""

1293 + cpu_offload_gb: float = 0

1294 + """The space in GiB to offload to CPU, per GPU. Default is 0, which means

1295 + no offloading. Intuitively, this argument can be seen as a virtual way to

1296 + increase the GPU memory size. For example, if you have one 24 GB GPU and

1297 + set this to 10, virtually you can think of it as a 34 GB GPU. Then you can

1298 + load a 13B model with BF16 weight, which requires at least 26GB GPU memory.

1299 + Note that this requires fast CPU-GPU interconnect, as part of the model is

1300 + loaded from CPU memory to GPU memory on the fly in each model forward pass.

1301 """

1302 + calculate_kv_scales: bool = False

1303 + """This enables dynamic calculation of `k_scale` and `v_scale` when

1304 + kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model

1305 + checkpoint if available. Otherwise, the scales will default to 1.0."""

1306 +

1307 + # Will be set after profiling.

1308 + num_gpu_blocks: Optional[int] = field(default=None, init=False)

1309 + """The number of blocks to allocate for GPU memory."""

1310 + num_cpu_blocks: Optional[int] = field(default=None, init=False)

1311 + """The number of blocks to allocate for CPU memory."""

1312

1313 def compute_hash(self) -> str:

1314 """

@@ -1281,43 +1329,13 @@ class CacheConfig:

1329 usedforsecurity=False).hexdigest()

1330 return hash_str

1331

1284 - def __init__(

1285 - self,

1286 - block_size: int,

1287 - gpu_memory_utilization: float,

1288 - swap_space: float,

1289 - cache_dtype: str,

1290 - is_attention_free: bool = False,

1291 - num_gpu_blocks_override: Optional[int] = None,

1292 - sliding_window: Optional[int] = None,

1293 - enable_prefix_caching: bool = False,

1294 - prefix_caching_hash_algo: str = "builtin",

1295 - cpu_offload_gb: float = 0,

1296 - calculate_kv_scales: Optional[bool] = None,

1297 - ) -> None:

1298 - self.block_size = block_size

1299 - self.gpu_memory_utilization = gpu_memory_utilization

1300 - self.swap_space_bytes = swap_space * GiB_bytes

1301 - self.num_gpu_blocks_override = num_gpu_blocks_override

1302 - self.cache_dtype = cache_dtype

1303 - self.is_attention_free = is_attention_free

1304 - self.sliding_window = sliding_window

1305 - self.enable_prefix_caching = enable_prefix_caching

1306 - self.prefix_caching_hash_algo = prefix_caching_hash_algo

1307 - self.cpu_offload_gb = cpu_offload_gb

1308 - self.calculate_kv_scales = calculate_kv_scales

1332 + def __post_init__(self) -> None:

1333 + self.swap_space_bytes = self.swap_space * GiB_bytes

1334 +

1335 self._verify_args()

1336 self._verify_cache_dtype()

1337 self._verify_prefix_caching()

1338

1313 - # Will be set after profiling.

1314 - self.num_gpu_blocks: Optional[int] = None

1315 - self.num_cpu_blocks: Optional[int] = None

1316 -

1317 - # Set calculate_kv_scales to False if the value is unset.

1318 - if self.calculate_kv_scales is None:

1319 - self.calculate_kv_scales = False

1320 -

1339 def metrics_info(self):

1340 # convert cache_config to dict(key: str, value: str) for prometheus

1341 # metrics info

@@ -1336,7 +1354,7 @@ class CacheConfig:

1354 def _verify_cache_dtype(self) -> None:

1355 if self.cache_dtype == "auto":

1356 pass

1339 - elif self.cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"):

1357 + elif self.cache_dtype in get_args(CacheDType):

1358 logger.info(

1359 "Using fp8 data type to store kv cache. It reduces the GPU "

1360 "memory footprint and boosts the performance. "

@@ -1354,12 +1372,12 @@ class CacheConfig:

1372 "Prefix caching is not supported with sliding window. "

1373 "Run with --disable-sliding-window to use prefix caching.")

1374

1357 - if self.enable_prefix_caching and self.prefix_caching_hash_algo not in (

1358 - "builtin", "sha256"):

1375 + if (self.enable_prefix_caching and self.prefix_caching_hash_algo

1376 + not in get_args(PrefixCachingHashAlgo)):

1377 raise ValueError(

1378 "Unknown prefix caching hash algorithm: "

1361 - f"{self.prefix_caching_hash_algo}. Must be either "

1362 - "'builtin' or 'sha256'.")

1379 + f"{self.prefix_caching_hash_algo}. Must be one of "

1380 + f"{get_args(PrefixCachingHashAlgo)}.")

1381

1382 def verify_with_parallel_config(

1383 self,

vllm/engine/arg_utils.py

+53 -105

@@ -16,16 +16,16 @@ from typing_extensions import TypeIs

17 import vllm.envs as envs

18 from vllm import version

19 -from vllm.config import (CacheConfig, CompilationConfig, Config, ConfigFormat,

20 - DecodingConfig, Device, DeviceConfig,

21 - DistributedExecutorBackend, HfOverrides,

19 +from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,

20 + Config, ConfigFormat, DecodingConfig, Device,

21 + DeviceConfig, DistributedExecutorBackend, HfOverrides,

22 KVTransferConfig, LoadConfig, LoadFormat, LoRAConfig,

23 ModelConfig, ModelImpl, MultiModalConfig,

24 ObservabilityConfig, ParallelConfig, PoolerConfig,

25 - PoolType, PromptAdapterConfig, SchedulerConfig,

26 - SchedulerPolicy, SpeculativeConfig, TaskOption,

27 - TokenizerPoolConfig, VllmConfig, get_attr_docs,

28 - get_field)

25 + PoolType, PrefixCachingHashAlgo, PromptAdapterConfig,

26 + SchedulerConfig, SchedulerPolicy, SpeculativeConfig,

27 + TaskOption, TokenizerPoolConfig, VllmConfig,

28 + get_attr_docs, get_field)

29 from vllm.executor.executor_base import ExecutorBase

30 from vllm.logger import init_logger

31 from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS

@@ -138,7 +138,7 @@ class EngineArgs:

138 load_format: str = LoadConfig.load_format

139 config_format: ConfigFormat = ConfigFormat.AUTO

140 dtype: str = 'auto'

141 - kv_cache_dtype: str = 'auto'

141 + kv_cache_dtype: CacheDType = CacheConfig.cache_dtype

142 seed: Optional[int] = None

143 max_model_len: Optional[int] = None

144 # Note: Specifying a custom executor backend by passing a class

@@ -154,15 +154,16 @@ class EngineArgs:

154 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel

155 max_parallel_loading_workers: Optional[

156 int] = ParallelConfig.max_parallel_loading_workers

157 - block_size: Optional[int] = None

158 - enable_prefix_caching: Optional[bool] = None

159 - prefix_caching_hash_algo: str = "builtin"

157 + block_size: Optional[BlockSize] = CacheConfig.block_size

158 + enable_prefix_caching: Optional[bool] = CacheConfig.enable_prefix_caching

159 + prefix_caching_hash_algo: PrefixCachingHashAlgo = \

160 + CacheConfig.prefix_caching_hash_algo

161 disable_sliding_window: bool = False

162 disable_cascade_attn: bool = False

163 use_v2_block_manager: bool = True

163 - swap_space: float = 4 # GiB

164 - cpu_offload_gb: float = 0 # GiB

165 - gpu_memory_utilization: float = 0.90

164 + swap_space: float = CacheConfig.swap_space

165 + cpu_offload_gb: float = CacheConfig.cpu_offload_gb

166 + gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization

167 max_num_batched_tokens: Optional[

168 int] = SchedulerConfig.max_num_batched_tokens

169 max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills

@@ -211,7 +212,8 @@ class EngineArgs:

212 num_scheduler_steps: int = SchedulerConfig.num_scheduler_steps

213 multi_step_stream_outputs: bool = SchedulerConfig.multi_step_stream_outputs

214 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight

214 - num_gpu_blocks_override: Optional[int] = None

215 + num_gpu_blocks_override: Optional[

216 + int] = CacheConfig.num_gpu_blocks_override

217 num_lookahead_slots: int = SchedulerConfig.num_lookahead_slots

218 model_loader_extra_config: dict = \

219 get_field(LoadConfig, "model_loader_extra_config")

@@ -250,7 +252,7 @@ class EngineArgs:

252 enable_sleep_mode: bool = False

253 model_impl: str = "auto"

254

253 - calculate_kv_scales: Optional[bool] = None

255 + calculate_kv_scales: bool = CacheConfig.calculate_kv_scales

256

257 additional_config: Optional[Dict[str, Any]] = None

258 enable_reasoning: Optional[bool] = None

@@ -306,12 +308,19 @@ class EngineArgs:

308 cls_docs = get_attr_docs(cls)

309 kwargs = {}

310 for field in fields(cls):

309 - name = field.name

311 + # Get the default value of the field

312 default = field.default

311 - # This will only be True if default is MISSING

313 if field.default_factory is not MISSING:

314 default = field.default_factory()

314 - kwargs[name] = {"default": default, "help": cls_docs[name]}

315 +

316 + # Get the help text for the field

317 + name = field.name

318 + help = cls_docs[name]

319 + # Escape % for argparse

320 + help = help.replace("%", "%%")

321 +

322 + # Initialise the kwargs dictionary for the field

323 + kwargs[name] = {"default": default, "help": help}

324

325 # Make note of if the field is optional and get the actual

326 # type of the field if it is

@@ -319,6 +328,8 @@ class EngineArgs:

328 field_type = get_args(

329 field.type)[0] if optional else field.type

330

331 + # Set type, action and choices for the field depending on the

332 + # type of the field

333 if can_be_type(field_type, bool):

334 # Creates --no-<name> and --<name> flags

335 kwargs[name]["action"] = argparse.BooleanOptionalAction

@@ -463,14 +474,6 @@ class EngineArgs:

474 '* "bfloat16" for a balance between precision and range.\n'

475 '* "float" is shorthand for FP32 precision.\n'

476 '* "float32" for FP32 precision.')

466 - parser.add_argument(

467 - '--kv-cache-dtype',

468 - type=str,

469 - choices=['auto', 'fp8', 'fp8_e5m2', 'fp8_e4m3'],

470 - default=EngineArgs.kv_cache_dtype,

471 - help='Data type for kv cache storage. If "auto", will use model '

472 - 'data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. '

473 - 'ROCm (AMD GPU) supports fp8 (=fp8_e4m3)')

477 parser.add_argument('--max-model-len',

478 type=human_readable_int,

479 default=EngineArgs.max_model_len,

@@ -544,33 +547,30 @@ class EngineArgs:

547 parallel_group.add_argument(

548 '--disable-custom-all-reduce',

549 **parallel_kwargs["disable_custom_all_reduce"])

547 - # KV cache arguments

548 - parser.add_argument('--block-size',

549 - type=int,

550 - default=EngineArgs.block_size,

551 - choices=[8, 16, 32, 64, 128],

552 - help='Token block size for contiguous chunks of '

553 - 'tokens. This is ignored on neuron devices and '

554 - 'set to ``--max-model-len``. On CUDA devices, '

555 - 'only block sizes up to 32 are supported. '

556 - 'On HPU devices, block size defaults to 128.')

550

558 - parser.add_argument(

559 - "--enable-prefix-caching",

560 - action=argparse.BooleanOptionalAction,

561 - default=EngineArgs.enable_prefix_caching,

562 - help="Enables automatic prefix caching. "

563 - "Use ``--no-enable-prefix-caching`` to disable explicitly.",

564 - )

565 - parser.add_argument(

566 - "--prefix-caching-hash-algo",

567 - type=str,

568 - choices=["builtin", "sha256"],

569 - default=EngineArgs.prefix_caching_hash_algo,

570 - help="Set the hash algorithm for prefix caching. "

571 - "Options are 'builtin' (Python's built-in hash) or 'sha256' "

572 - "(collision resistant but with certain overheads).",

551 + # KV cache arguments

552 + cache_kwargs = get_kwargs(CacheConfig)

553 + cache_group = parser.add_argument_group(

554 + title="CacheConfig",

555 + description=CacheConfig.__doc__,

556 )

557 + cache_group.add_argument('--block-size', **cache_kwargs["block_size"])

558 + cache_group.add_argument('--gpu-memory-utilization',

559 + **cache_kwargs["gpu_memory_utilization"])

560 + cache_group.add_argument('--swap-space', **cache_kwargs["swap_space"])

561 + cache_group.add_argument('--kv-cache-dtype',

562 + **cache_kwargs["cache_dtype"])

563 + cache_group.add_argument('--num-gpu-blocks-override',

564 + **cache_kwargs["num_gpu_blocks_override"])

565 + cache_group.add_argument("--enable-prefix-caching",

566 + **cache_kwargs["enable_prefix_caching"])

567 + cache_group.add_argument("--prefix-caching-hash-algo",

568 + **cache_kwargs["prefix_caching_hash_algo"])

569 + cache_group.add_argument('--cpu-offload-gb',

570 + **cache_kwargs["cpu_offload_gb"])

571 + cache_group.add_argument('--calculate-kv-scales',

572 + **cache_kwargs["calculate_kv_scales"])

573 +

574 parser.add_argument('--disable-sliding-window',

575 action='store_true',

576 help='Disables sliding window, '

@@ -588,43 +588,6 @@ class EngineArgs:

588 type=int,

589 default=EngineArgs.seed,

590 help='Random seed for operations.')

591 - parser.add_argument('--swap-space',

592 - type=float,

593 - default=EngineArgs.swap_space,

594 - help='CPU swap space size (GiB) per GPU.')

595 - parser.add_argument(

596 - '--cpu-offload-gb',

597 - type=float,

598 - default=0,

599 - help='The space in GiB to offload to CPU, per GPU. '

600 - 'Default is 0, which means no offloading. Intuitively, '

601 - 'this argument can be seen as a virtual way to increase '

602 - 'the GPU memory size. For example, if you have one 24 GB '

603 - 'GPU and set this to 10, virtually you can think of it as '

604 - 'a 34 GB GPU. Then you can load a 13B model with BF16 weight, '

605 - 'which requires at least 26GB GPU memory. Note that this '

606 - 'requires fast CPU-GPU interconnect, as part of the model is '

607 - 'loaded from CPU memory to GPU memory on the fly in each '

608 - 'model forward pass.')

609 - parser.add_argument(

610 - '--gpu-memory-utilization',

611 - type=float,

612 - default=EngineArgs.gpu_memory_utilization,

613 - help='The fraction of GPU memory to be used for the model '

614 - 'executor, which can range from 0 to 1. For example, a value of '

615 - '0.5 would imply 50%% GPU memory utilization. If unspecified, '

616 - 'will use the default value of 0.9. This is a per-instance '

617 - 'limit, and only applies to the current vLLM instance.'

618 - 'It does not matter if you have another vLLM instance running '

619 - 'on the same GPU. For example, if you have two vLLM instances '

620 - 'running on the same GPU, you can set the GPU memory utilization '

621 - 'to 0.5 for each instance.')

622 - parser.add_argument(

623 - '--num-gpu-blocks-override',

624 - type=int,

625 - default=None,

626 - help='If specified, ignore GPU profiling result and use this number'

627 - ' of GPU blocks. Used for testing preemption.')

591 parser.add_argument(

592 '--max-logprobs',

593 type=int,

@@ -994,15 +957,6 @@ class EngineArgs:

957 help="Enable sleep mode for the engine. "

958 "(only cuda platform is supported)")

959

997 - parser.add_argument(

998 - '--calculate-kv-scales',

999 - action='store_true',

1000 - help='This enables dynamic calculation of '

1001 - 'k_scale and v_scale when kv-cache-dtype is fp8. '

1002 - 'If calculate-kv-scales is false, the scales will '

1003 - 'be loaded from the model checkpoint if available. '

1004 - 'Otherwise, the scales will default to 1.0.')

1005 -

960 parser.add_argument(

961 "--additional-config",

962 type=json.loads,

@@ -1625,9 +1579,7 @@ class EngineArgs:

1579 self.enable_prefix_caching = False

1580

1581 # VLLM_V0 only supports builtin hash algo for prefix caching.

1628 - if self.prefix_caching_hash_algo is None:

1629 - self.prefix_caching_hash_algo = "builtin"

1630 - elif self.prefix_caching_hash_algo == "sha256":

1582 + if self.prefix_caching_hash_algo == "sha256":

1583 raise ValueError(

1584 "sha256 is not supported for prefix caching in V0 engine. "

1585 "Please use 'builtin'.")

@@ -1646,10 +1598,6 @@ class EngineArgs:

1598 if self.enable_prefix_caching is None:

1599 self.enable_prefix_caching = True

1600

1649 - # if using prefix caching, we must set a hash algo

1650 - if self.enable_prefix_caching and self.prefix_caching_hash_algo is None:

1651 - self.prefix_caching_hash_algo = "builtin"

1652 -

1601 # V1 should use the new scheduler by default.

1602 # Swap it only if this arg is set to the original V0 default

1603 if self.scheduler_cls == EngineArgs.scheduler_cls:

vllm/platforms/neuron.py

+1 -1

@@ -50,7 +50,7 @@ class NeuronPlatform(Platform):

50 if cache_config:

51 # neuron needs block_size = max_model_len

52 vllm_config.cache_config.block_size = \

53 - vllm_config.model_config.max_model_len

53 + vllm_config.model_config.max_model_len # type: ignore

55 @classmethod

56 def is_pin_memory_available(cls) -> bool: