@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4b1d141f49
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, 21:00 UTC

[PP] Correct cache size check (#13873)

Signed-off-by: Yang Zheng <zhengy.gator@gmail.com>

Yang Zheng committed Feb 27, 2025 at 17:47 UTC 4b1d141f49559d1efb36d2702ebce906fd414cb2
2 files changed +14 -12
vllm/worker/hpu_worker.py
+7 -6
@@ -258,9 +258,10 @@ class HPUWorker(LocalOrDistributedWorkerBase):
258
259 This also warms up the model, which may record CUDA graphs.
260 """
261 - raise_if_cache_size_invalid(num_gpu_blocks,
262 - self.cache_config.block_size,
263 - self.model_config.max_model_len)
261 + raise_if_cache_size_invalid(
262 + num_gpu_blocks, self.cache_config.block_size,
263 + self.model_config.max_model_len,
264 + self.parallel_config.pipeline_parallel_size)
265
266 self.cache_config.num_gpu_blocks = num_gpu_blocks
267 self.cache_config.num_cpu_blocks = num_cpu_blocks
@@ -442,13 +443,13 @@ def init_worker_distributed_environment(
443 parallel_config.pipeline_parallel_size)
444
445
445 -def raise_if_cache_size_invalid(num_gpu_blocks, block_size,
446 - max_model_len) -> None:
446 +def raise_if_cache_size_invalid(num_gpu_blocks, block_size, max_model_len,
447 + pipeline_parallel_size) -> None:
448 if num_gpu_blocks <= 0:
449 raise ValueError("No available memory for the cache blocks. "
450 "Try increasing `gpu_memory_utilization` when "
451 "initializing the engine.")
451 - max_seq_len = block_size * num_gpu_blocks
452 + max_seq_len = block_size * (num_gpu_blocks // pipeline_parallel_size)
453 if max_model_len > max_seq_len:
454 raise ValueError(
455 f"The model's max seq len ({max_model_len}) "
vllm/worker/worker.py
+7 -6
@@ -288,10 +288,11 @@ class Worker(LocalOrDistributedWorkerBase):
288
289 This also warms up the model, which may record CUDA graphs.
290 """
291 - raise_if_cache_size_invalid(num_gpu_blocks,
292 - self.cache_config.block_size,
293 - self.cache_config.is_attention_free,
294 - self.model_config.max_model_len)
291 + raise_if_cache_size_invalid(
292 + num_gpu_blocks, self.cache_config.block_size,
293 + self.cache_config.is_attention_free,
294 + self.model_config.max_model_len,
295 + self.parallel_config.pipeline_parallel_size)
296
297 self.cache_config.num_gpu_blocks = num_gpu_blocks
298 self.cache_config.num_cpu_blocks = num_cpu_blocks
@@ -530,7 +531,7 @@ def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
531
532
533 def raise_if_cache_size_invalid(num_gpu_blocks, block_size, is_attention_free,
533 - max_model_len) -> None:
534 + max_model_len, pipeline_parallel_size) -> None:
535 if is_attention_free and num_gpu_blocks != 0:
536 raise ValueError("No memory should be allocated for the cache blocks "
537 f"for an attention-free model, but {num_gpu_blocks} "
@@ -539,7 +540,7 @@ def raise_if_cache_size_invalid(num_gpu_blocks, block_size, is_attention_free,
540 raise ValueError("No available memory for the cache blocks. "
541 "Try increasing `gpu_memory_utilization` when "
542 "initializing the engine.")
542 - max_seq_len = block_size * num_gpu_blocks
543 + max_seq_len = block_size * (num_gpu_blocks // pipeline_parallel_size)
544 if not is_attention_free and max_model_len > max_seq_len:
545 raise ValueError(
546 f"The model's max seq len ({max_model_len}) "