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, 12:15 UTC
[Misc]add configurable cuda graph size (#17201)
Signed-off-by: CXIAAAAA <cxia0209@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Chen Xia committed
May 1, 2025 at 11:04 UTC
61c299f81fdf4df30f699a5030f1eb4b10d0cdba
2 files changed
+22
-3
vllm/config.py
+17
-3
@@ -1865,6 +1865,13 @@ class SchedulerConfig:
1865
This config has no static default. If left unspecified by the user, it will
1866
be set in `EngineArgs.create_engine_config` based on the usage context."""
1867
1868
+ cuda_graph_sizes: list[int] = field(default_factory=lambda: [512])
1869
+ """Cuda graph capture sizes, default is 512.
1870
+ 1. if one value is provided, then the capture list would follow the pattern:
1871
+ [1, 2, 4] + [i for i in range(8, cuda_graph_sizes + 1, 8)]
1872
+ 2. more than one value (e.g. 1 2 128) is provided,
1873
+ then the capture list will follow the provided list."""
1874
+
1875
max_num_seqs: int = None # type: ignore
1876
"""Maximum number of sequences to be processed in a single iteration.
1877
@@ -4235,13 +4242,20 @@ class VllmConfig:
4242
batch_size_capture_list = []
4243
if self.model_config is not None and \
4244
not self.model_config.enforce_eager:
4238
- batch_size_capture_list = [1, 2, 4
4239
- ] + [i for i in range(8, 513, 8)]
4245
+ cuda_graph_sizes = self.scheduler_config.cuda_graph_sizes
4246
+ if len(cuda_graph_sizes) == 1:
4247
+ batch_size_capture_list = [1, 2, 4] + [
4248
+ i for i in range(8, cuda_graph_sizes[0] + 1, 8)
4249
+ ]
4250
+ elif len(cuda_graph_sizes) > 1:
4251
+ batch_size_capture_list = sorted(cuda_graph_sizes)
4252
+ else:
4253
+ raise TypeError(
4254
+ f"Invalid value for {cuda_graph_sizes=}.")
4255
if self.parallel_config.tensor_parallel_size > 1 and \
4256
self.compilation_config.pass_config.enable_sequence_parallelism:
4257
batch_size_capture_list = \
4258
self.update_sizes_for_sequence_parallelism(batch_size_capture_list)
4244
-
4259
max_num_tokens = self.scheduler_config.max_num_batched_tokens
4260
batch_size_capture_list = [
4261
size for size in batch_size_capture_list
vllm/engine/arg_utils.py
+5
@@ -231,6 +231,8 @@ class EngineArgs:
231
kv_cache_dtype: CacheDType = CacheConfig.cache_dtype
232
seed: Optional[int] = ModelConfig.seed
233
max_model_len: Optional[int] = ModelConfig.max_model_len
234
+ cuda_graph_sizes: list[int] = get_field(SchedulerConfig,
235
+ "cuda_graph_sizes")
236
# Note: Specifying a custom executor backend by passing a class
237
# is intended for expert use only. The API may change without
238
# notice.
@@ -711,6 +713,8 @@ class EngineArgs:
713
scheduler_group.add_argument(
714
"--max-long-partial-prefills",
715
**scheduler_kwargs["max_long_partial_prefills"])
716
+ scheduler_group.add_argument('--cuda-graph-sizes',
717
+ **scheduler_kwargs["cuda_graph_sizes"])
718
scheduler_group.add_argument(
719
"--long-prefill-token-threshold",
720
**scheduler_kwargs["long_prefill_token_threshold"])
@@ -1042,6 +1046,7 @@ class EngineArgs:
1046
max_num_batched_tokens=self.max_num_batched_tokens,
1047
max_num_seqs=self.max_num_seqs,
1048
max_model_len=model_config.max_model_len,
1049
+ cuda_graph_sizes=self.cuda_graph_sizes,
1050
num_lookahead_slots=num_lookahead_slots,
1051
delay_factor=self.scheduler_delay_factor,
1052
enable_chunked_prefill=self.enable_chunked_prefill,