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

[torch.compile] Make inductor partition rules respect splitting_ops #25691 (#25845)

Browse Source 
Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com>
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
This commit is contained in:
baonudesifeizhai
2025-10-10 12:35:28 -04:00
committed by GitHub
parent e519281920
commit cddce79fda
9 changed files with 267 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -127,7 +127,7 @@ def _run_simple_model(
@torch.inference_mode()
def test_simple_piecewise_compile(use_inductor):
_run_simple_model(
splitting_ops=["silly.attention"],
splitting_ops=["silly::attention"],
use_inductor_graph_partition=False,
use_inductor=use_inductor,
# 2 * num_layers + 1
@@ -142,7 +142,7 @@ def test_simple_piecewise_compile(use_inductor):
@torch.inference_mode()
@pytest.mark.parametrize("splitting_ops", [["silly.attention"], []])
@pytest.mark.parametrize("splitting_ops", [["silly::attention"], []])
def test_simple_inductor_graph_partition(splitting_ops, monkeypatch):
if not is_torch_equal_or_newer("2.9.0.dev"):
pytest.skip("inductor graph partition is only available in PyTorch 2.9+")