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 14, 05:30 UTC
[Feature] Add `VLLM_DISABLE_PAD_FOR_CUDAGRAPH` to Avoid Hang Issue (#23595)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Wentao Ye committed
Aug 27, 2025 at 17:52 UTC
321938e9ac4000e0cb37e328359a7fd3026bc672
2 files changed
+8
vllm/envs.py
+7
@@ -166,6 +166,7 @@ if TYPE_CHECKING:
166
VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
167
VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False
168
VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
169
+ VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False
170
171
172
def get_default_cache_root():
@@ -1144,6 +1145,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1145
"VLLM_ENABLE_CUDAGRAPH_GC":
1146
lambda: bool(int(os.getenv("VLLM_ENABLE_CUDAGRAPH_GC", "0"))),
1147
1148
+ # Disable padding to CUDA graph capture batch sizes.
1149
+ # TODO(wentao): https://github.com/vllm-project/vllm/issues/23378
1150
+ # After the issue is fixed, we can remove this flag.
1151
+ "VLLM_DISABLE_PAD_FOR_CUDAGRAPH":
1152
+ lambda: bool(int(os.getenv("VLLM_DISABLE_PAD_FOR_CUDAGRAPH", "0"))),
1153
+
1154
# Used to force set up loopback IP
1155
"VLLM_LOOPBACK_IP":
1156
lambda: os.getenv("VLLM_LOOPBACK_IP", ""),
vllm/v1/worker/gpu_model_runner.py
+1
@@ -1491,6 +1491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1491
1492
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1493
if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
1494
+ and not envs.VLLM_DISABLE_PAD_FOR_CUDAGRAPH
1495
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
1496
# Use CUDA graphs.
1497
# Add padding to the batch size.