Skip to content

Commit

Permalink
[torch.compile] PyTorch 2.6 and nightly compatibility (vllm-project#1…
Browse files Browse the repository at this point in the history
…2393)

Signed-off-by: youkaichao <youkaichao@gmail.com>
  • Loading branch information
youkaichao authored and AoyuQC committed Feb 8, 2025
1 parent d454c29 commit efbe885
Show file tree
Hide file tree
Showing 8 changed files with 493 additions and 320 deletions.
2 changes: 1 addition & 1 deletion tests/compile/piecewise/test_simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ def test_simple_piecewise_compile():
num_graphs_seen=1, # one graph for the model
num_piecewise_graphs_seen=5, # 2 * num_layers + 1
num_piecewise_capturable_graphs_seen=3, # 1 + num_layers
num_inductor_compilations=3, # num_piecewise_capturable_graphs_seen
num_backend_compilations=3, # num_piecewise_capturable_graphs_seen
num_cudagraph_caputured=
6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
):
Expand Down
6 changes: 3 additions & 3 deletions tests/compile/piecewise/test_toy_llama.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ def test_toy_llama():
num_graphs_seen=0,
num_piecewise_graphs_seen=0,
num_piecewise_capturable_graphs_seen=0,
num_inductor_compilations=0,
num_backend_compilations=0,
num_cudagraph_caputured=0,
):
outputs.append(run_model(llama_config, use_compile=False))
Expand All @@ -332,7 +332,7 @@ def test_toy_llama():
num_graphs_seen=1, # one graph for the model
num_piecewise_graphs_seen=1,
num_piecewise_capturable_graphs_seen=1,
num_inductor_compilations=1, # num_piecewise_capturable_graphs_seen
num_backend_compilations=1, # num_piecewise_capturable_graphs_seen
num_cudagraph_caputured=
2, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
):
Expand All @@ -345,7 +345,7 @@ def test_toy_llama():
1, # 2 * num_layers + 1
num_piecewise_capturable_graphs_seen=1 +
llama_config.num_layers, # 1 + num_layers
num_inductor_compilations=1 +
num_backend_compilations=1 +
llama_config.num_layers, # num_piecewise_capturable_graphs_seen
num_cudagraph_caputured=2 *
(1 + llama_config.num_layers
Expand Down
Loading

0 comments on commit efbe885

Please sign in to comment.