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

Signed-off-by: youkaichao <youkaichao@gmail.com>
2024-11-20 11:20:38 -08:00
parent 5f1d6af2b6
commit 0cd3d9717e
5 changed files with 27 additions and 16 deletions
--- a/tests/compile/piecewise/test_simple.py
+++ b/tests/compile/piecewise/test_simple.py
@@ -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):