"vscode:/vscode.git/clone" did not exist on "4df841fe7538cb8de281b9d78e37ba51ac35b5da"
test_simple.py 3.35 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
"""
Test the piecewise compilation with a simple model so that we
can exactly calculate the expected output and side effects.
"""

import torch
from torch import nn
10
from torch.library import Library
11
12
13

from vllm.compilation.counter import compilation_counter
from vllm.compilation.decorators import support_torch_compile
14
15
from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
                         set_current_vllm_config)
16
from vllm.utils import direct_register_custom_op
17
18
19

global_counter = 0

20
21
22
# create a library to hold the custom op
silly_lib = Library("silly", "FRAGMENT")  # noqa

23
24
25
26
27
28
29
30
31
32

def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                    out: torch.Tensor) -> None:
    global global_counter
    global_counter += 1
    print(f"{global_counter=}")
    out.copy_(q)
    out[0] += 1


33
34
def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                         out: torch.Tensor) -> None:
35
36
37
    return


38
39
40
41
42
43
44
45
46
direct_register_custom_op(
    op_name="attention",
    op_func=silly_attention,
    mutates_args=["out"],
    fake_impl=silly_attention_fake,
    target_lib=silly_lib,
)


47
48
49
@support_torch_compile
class SillyModel(nn.Module):

50
51
52
53
54
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 prefix: str = '',
                 **kwargs) -> None:
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Overall effect:
        x += 1
        x[0] += 2
        global_counter += 2
        """
        x = x + 1
        x = x + 2
        out = torch.empty_like(x)
        torch.ops.silly.attention(x, x, x, out)
        x = out
        x = x - 2
        x = x - 1
        out = torch.empty_like(x)
        torch.ops.silly.attention(x, x, x, out)
        x = out
        x = x + 1
        return x


78
def _test_simple_piecewise_compile(*, use_inductor):
79

80
81
82
    vllm_config = VllmConfig(compilation_config=CompilationConfig(
        level=CompilationLevel.PIECEWISE,
        use_cudagraph=True,
83
        use_inductor=use_inductor,
84
        splitting_ops=["silly.attention"],
85
        cudagraph_copy_inputs=True,
86
        cudagraph_capture_sizes=[1, 2],
87
    ))
88
89
    with set_current_vllm_config(vllm_config):
        model = SillyModel(vllm_config=vllm_config, prefix='')
90

91
    inputs = torch.randn(100).cuda()
92
93
94
95
96

    with compilation_counter.expect(
            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
97
            num_backend_compilations=3,  # num_piecewise_capturable_graphs_seen
98
99
100
101
            num_cudagraph_caputured=
            6,  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
    ):

102
        model(inputs)
103

104
105
        model(torch.randn(2).cuda())
        model(torch.randn(1).cuda())
106

107
        input = torch.zeros(2).cuda()
108
109
        global global_counter
        global_counter = 0
110
        output = model(input)
111
112
        assert global_counter == 2
        assert torch.allclose(output.cpu(), torch.tensor([3., 1.]))
113
114
115
116
117
118
119
120


def test_simple_piecewise_compile_inductor():
    _test_simple_piecewise_compile(use_inductor=True)


def test_simple_piecewise_compile_no_inductor():
    _test_simple_piecewise_compile(use_inductor=False)