[V1] Enable custom ops with piecewise CUDA graphs (#10228)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2024-11-11 11:58:07 -08:00 committed by GitHub
parent 8a7fe47d32
commit 9d5b4e4dea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,3 +1,4 @@
import os
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Set
@ -405,6 +406,7 @@ class GPUModelRunner:
if self.use_cuda_graph:
# FIXME(woosuk): Currently, we do not use inductor to reduce the
# compilation time and any potential issues with the inductor.
os.environ["VLLM_CUSTOM_OPS"] = "all"
set_compilation_config(
CompilationConfig(
use_cudagraph=True,