youngkingdom/vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[7/N] torch.compile, reduce compilation time (#10460)

Browse Source 
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-11-20 11:20:38 -08:00
committed by GitHub
parent 5f1d6af2b6
commit 0cd3d9717e
5 changed files with 27 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/compile/piecewise/test_simple.py
View File

@ -79,7 +79,7 @@ def test_simple_piecewise_compile():
vllm_config = VllmConfig(compilation_config=CompilationConfig(
level=CompilationLevel.PIECEWISE,
use_cudagraph=True,
non_cudagraph_ops=["silly.attention"],
splitting_ops=["silly.attention"],
cudagraph_copy_inputs=True,
))
with set_current_vllm_config(vllm_config):

4
tests/compile/piecewise/test_toy_llama.py
View File

@ -258,7 +258,7 @@ def run_model(llama_config,
use_cudagraph=True,
)
if split_attn:
compilation_config.non_cudagraph_ops = ["silly.attention"]
compilation_config.splitting_ops = ["silly.attention"]
else:
compilation_config = CompilationConfig(
level=CompilationLevel.NO_COMPILATION, )
@ -378,7 +378,7 @@ def benchmark():
compilation_config = CompilationConfig(
level=CompilationLevel.PIECEWISE,
use_cudagraph=True,
non_cudagraph_ops=["silly.attention"],
splitting_ops=["silly.attention"],
)
else:
compilation_config = CompilationConfig(