Skip to content

GitLab

Unverified Commit 9d5b4e4d authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

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


Signed-off-by: default avatarWoosuk Kwon <woosuk.kwon@berkeley.edu>
parent 8a7fe47d
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
vllm/v1/worker/gpu_model_runner.py
View file @ 9d5b4e4d
import os
import time import time
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Set from typing import TYPE_CHECKING, Dict, List, Optional, Set
...@@ -405,6 +406,7 @@ class GPUModelRunner: ...@@ -405,6 +406,7 @@ class GPUModelRunner:
if self.use_cuda_graph: if self.use_cuda_graph:
# FIXME(woosuk): Currently, we do not use inductor to reduce the # FIXME(woosuk): Currently, we do not use inductor to reduce the
# compilation time and any potential issues with the inductor. # compilation time and any potential issues with the inductor.
os.environ["VLLM_CUSTOM_OPS"] = "all"
set_compilation_config( set_compilation_config(
CompilationConfig( CompilationConfig(
use_cudagraph=True, use_cudagraph=True,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment