A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI] execute all piecewise compilation tests together (#24502)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Sep 10, 2025 at 02:05 UTC
b8a93076d36eff5cff8a89f99a7370d0cc6f0e98
6 files changed
+81
-117
.buildkite/test-pipeline.yaml
+1
-5
@@ -379,11 +379,7 @@ steps:
379
- tests/compile
380
commands:
381
- pytest -v -s compile/test_basic_correctness.py
382
- # these tests need to be separated, cannot combine
383
- - pytest -v -s compile/piecewise/test_simple.py
384
- - pytest -v -s compile/piecewise/test_toy_llama.py
385
- - pytest -v -s compile/piecewise/test_full_cudagraph.py
386
- - pytest -v -s compile/piecewise/test_multiple_graphs.py
382
+ - pytest -v -s compile/piecewise/
383
384
- label: PyTorch Fullgraph Test # 20min
385
timeout_in_minutes: 30
tests/compile/piecewise/test_multiple_graphs.py
+3
-25
@@ -4,9 +4,9 @@
4
Test (piecewise) compilation with a simple model where multiple submodules
5
are compiled and graph captured separately.
6
"""
7
+
8
import torch
9
from torch import nn
9
-from torch.library import Library
10
11
from vllm.compilation.backends import set_model_tag
12
from vllm.compilation.counter import compilation_counter
@@ -15,10 +15,9 @@ from vllm.compilation.decorators import (ignore_torch_compile,
15
from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
16
VllmConfig, set_current_vllm_config)
17
from vllm.forward_context import BatchDescriptor, set_forward_context
18
-from vllm.utils import direct_register_custom_op
18
20
-# create a library to hold the custom op
21
-silly_lib = Library("silly", "FRAGMENT") # noqa
19
+# This import automatically registers `torch.ops.silly.attention`
20
+from .. import silly_attention # noqa: F401
21
22
BATCH_SIZE = 32
23
MLP_SIZE = 128
@@ -26,27 +25,6 @@ HIDDEN_SIZE = 1024
25
RANDOM_SEED = 0
26
27
29
-def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
30
- out: torch.Tensor) -> None:
31
- out.copy_(q)
32
- out += k
33
- out += v
34
-
35
-
36
-def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
37
- out: torch.Tensor) -> None:
38
- return
39
-
40
-
41
-direct_register_custom_op(
42
- op_name="attention",
43
- op_func=silly_attention,
44
- mutates_args=["out"],
45
- fake_impl=silly_attention_fake,
46
- target_lib=silly_lib,
47
-)
48
-
49
-
28
@support_torch_compile
29
class ParentModel(nn.Module):
30
tests/compile/piecewise/test_simple.py
+8
-35
@@ -4,10 +4,10 @@
4
Test the piecewise compilation with a simple model so that we
5
can exactly calculate the expected output and side effects.
6
"""
7
+
8
import pytest
9
import torch
10
from torch import nn
10
-from torch.library import Library
11
12
from vllm.compilation.counter import compilation_counter
13
from vllm.compilation.decorators import support_torch_compile
@@ -15,35 +15,9 @@ from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
15
VllmConfig, set_current_vllm_config)
16
from vllm.envs import VLLM_USE_V1
17
from vllm.forward_context import BatchDescriptor, set_forward_context
18
-from vllm.utils import direct_register_custom_op
19
-
20
-global_counter = 0
21
-
22
-# create a library to hold the custom op
23
-silly_lib = Library("silly", "FRAGMENT") # noqa
24
-
25
-
26
-def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
27
- out: torch.Tensor) -> None:
28
- global global_counter
29
- global_counter += 1
30
- print(f"{global_counter=}")
31
- out.copy_(q)
32
- out[0] += 1
33
-
34
-
35
-def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
36
- out: torch.Tensor) -> None:
37
- return
38
-
18
40
-direct_register_custom_op(
41
- op_name="attention",
42
- op_func=silly_attention,
43
- mutates_args=["out"],
44
- fake_impl=silly_attention_fake,
45
- target_lib=silly_lib,
46
-)
19
+# This import automatically registers `torch.ops.silly.attention`
20
+from ..silly_attention import get_global_counter, reset_global_counter
21
22
23
@support_torch_compile
@@ -59,8 +33,7 @@ class SillyModel(nn.Module):
33
def forward(self, x: torch.Tensor) -> torch.Tensor:
34
"""
35
Overall effect:
62
- x += 1
63
- x[0] += 2
36
+ x = 3 * x + 19
37
global_counter += 2
38
"""
39
x = x + 1
@@ -78,6 +51,7 @@ class SillyModel(nn.Module):
51
52
53
@pytest.mark.parametrize("use_inductor", [True, False])
54
+@torch.inference_mode()
55
def test_simple_piecewise_compile(use_inductor):
56
assert VLLM_USE_V1
57
@@ -121,13 +95,12 @@ def test_simple_piecewise_compile(use_inductor):
95
model(torch.randn(1).cuda())
96
97
input = torch.zeros(2).cuda()
124
- global global_counter
125
- global_counter = 0
98
+ reset_global_counter()
99
with set_forward_context(
100
None,
101
vllm_config=vllm_config,
102
cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
103
batch_descriptor=BatchDescriptor(num_tokens=2, )):
104
output = model(input)
132
- assert global_counter == 2
133
- assert torch.allclose(output.cpu(), torch.tensor([3., 1.]))
105
+ assert get_global_counter() == 2
106
+ assert torch.allclose(output.cpu(), torch.tensor([19.0, 19.0]))
tests/compile/piecewise/test_toy_llama.py
+2
-25
@@ -14,38 +14,15 @@ from typing import Any, Optional
14
import pytest
15
import torch
16
from torch import nn
17
-from torch.library import Library
17
18
from vllm.compilation.counter import compilation_counter
19
from vllm.compilation.decorators import support_torch_compile
20
from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
21
VllmConfig, set_current_vllm_config)
22
from vllm.forward_context import BatchDescriptor, set_forward_context
24
-from vllm.utils import direct_register_custom_op
23
26
-# create a library to hold the custom op
27
-silly_lib = Library("silly", "FRAGMENT") # noqa
28
-
29
-
30
-def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
31
- out: torch.Tensor) -> None:
32
- out.copy_(q)
33
- out += k
34
- out += v
35
-
36
-
37
-def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
38
- out: torch.Tensor) -> None:
39
- return
40
-
41
-
42
-direct_register_custom_op(
43
- op_name="attention",
44
- op_func=silly_attention,
45
- mutates_args=["out"],
46
- fake_impl=silly_attention_fake,
47
- target_lib=silly_lib,
48
-)
24
+# This import automatically registers `torch.ops.silly.attention`
25
+from .. import silly_attention # noqa: F401
26
27
28
@dataclass
tests/compile/silly_attention.py
new
+63
@@ -0,0 +1,63 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Shared PyTorch custom silly attention for compilation tests.
5
+Centralizes custom operation definitions to avoid duplicate registrations.
6
+"""
7
+
8
+import torch
9
+from torch.library import Library
10
+
11
+from vllm.utils import direct_register_custom_op
12
+
13
+# Shared library for all compilation test operations
14
+# Using "silly" namespace to match existing test expectations
15
+# import this file will automatically register
16
+# torch ops for testing (like silly.attention)
17
+silly_lib = Library("silly", "FRAGMENT")
18
+
19
+# Global counter that counts the number of times attention is invoked
20
+_global_counter = 0
21
+
22
+
23
+def get_global_counter():
24
+ """Get the current global counter value"""
25
+ return _global_counter
26
+
27
+
28
+def reset_global_counter():
29
+ """Reset the global counter to 0"""
30
+ global _global_counter
31
+ _global_counter = 0
32
+
33
+
34
+def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
35
+ out: torch.Tensor) -> None:
36
+ """
37
+ Unified attention implementation that depends on
38
+ all inputs and affects the output.
39
+ Always increments a global counter that tests can use or ignore.
40
+ """
41
+ global _global_counter
42
+
43
+ # Always increment the global counter
44
+ _global_counter += 1
45
+
46
+ # Unified implementation that depends on all inputs
47
+ out.copy_(q + k + v)
48
+
49
+
50
+def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
51
+ out: torch.Tensor) -> None:
52
+ """Fake implementation for testing"""
53
+ return
54
+
55
+
56
+# Register the unified attention operation
57
+direct_register_custom_op(
58
+ op_name="attention",
59
+ op_func=silly_attention,
60
+ mutates_args=["out"],
61
+ fake_impl=silly_attention_fake,
62
+ target_lib=silly_lib,
63
+)
tests/compile/test_decorator.py
+4
-27
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
from torch import nn
5
-from torch.library import Library
5
6
from vllm.compilation.counter import compilation_counter
7
from vllm.compilation.decorators import (ignore_torch_compile,
@@ -10,36 +9,14 @@ from vllm.compilation.decorators import (ignore_torch_compile,
9
from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
10
CUDAGraphMode, VllmConfig, set_current_vllm_config)
11
from vllm.forward_context import BatchDescriptor, set_forward_context
13
-from vllm.utils import direct_register_custom_op
12
15
-# create a library to hold the custom op
16
-silly_lib = Library("silly", "FRAGMENT") # noqa
13
+# This import automatically registers `torch.ops.silly.attention`
14
+from . import silly_attention # noqa: F401
15
16
BATCH_SIZE = 32
17
MLP_SIZE = 128
18
19
22
-def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
23
- out: torch.Tensor) -> None:
24
- out.copy_(q)
25
- out += k
26
- out += v
27
-
28
-
29
-def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
30
- out: torch.Tensor) -> None:
31
- return
32
-
33
-
34
-direct_register_custom_op(
35
- op_name="attention",
36
- op_func=silly_attention,
37
- mutates_args=["out"],
38
- fake_impl=silly_attention_fake,
39
- target_lib=silly_lib,
40
-)
41
-
42
-
20
@torch.inference_mode
21
def run_model(vllm_config: VllmConfig, model: nn.Module,
22
cudagraph_runtime_mode: CUDAGraphMode):
@@ -151,7 +128,7 @@ def test_ignore_torch_compile_decorator():
128
run_model(vllm_config, mod_C, cudagraph_runtime_mode)
129
130
154
-# Only enable torch.compile if
131
+# Only enable torch.compile if
132
# vllm_config.cache_config.kv_sharing_fast_prefill=True
133
@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.
134
kv_sharing_fast_prefill)
@@ -173,7 +150,7 @@ class B(nn.Module):
150
return x
151
152
176
-# Only enable torch.compile if
153
+# Only enable torch.compile if
154
# vllm_config.cache_config.kv_sharing_fast_prefill=False
155
@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.
156
cache_config.kv_sharing_fast_prefill)