Skip to content

Commit

Permalink
Add CUDA block size check
Browse files Browse the repository at this point in the history
  • Loading branch information
kzawora-intel committed Dec 17, 2024
1 parent 56a150e commit 007a573
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 1 deletion.
5 changes: 5 additions & 0 deletions vllm/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -917,6 +917,11 @@ def _verify_args(self) -> None:
raise ValueError(
"GPU memory utilization must be less than 1.0. Got "
f"{self.gpu_memory_utilization}.")
if (current_platform.is_cuda_alike()
and self.block_size not in [8, 16, 32]):
raise ValueError(
"CUDA Paged Attention kernel only supports "
f"block sizes [8, 16, 32]. Got {self.block_size}.")

def _verify_cache_dtype(self) -> None:
if self.cache_dtype == "auto":
Expand Down
4 changes: 3 additions & 1 deletion vllm/engine/arg_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
choices=[8, 16, 32, 64, 128],
help='Token block size for contiguous chunks of '
'tokens. This is ignored on neuron devices and '
'set to max-model-len')
'set to max-model-len. On CUDA devices, '
'only block sizes up to 32 are supported. '
'On HPU devices, block size defaults to 128.')

parser.add_argument(
"--enable-prefix-caching",
Expand Down

0 comments on commit 007a573

Please sign in to comment.