swiglu.py 6.86 KB
Newer Older
xgqdut2016's avatar
xgqdut2016 committed
1
import torch
PanZezhongQY's avatar
PanZezhongQY committed
2
import ctypes
3
from ctypes import POINTER, Structure, c_int32, c_void_p, c_uint64
xgqdut2016's avatar
xgqdut2016 committed
4
from libinfiniop import (
PanZezhongQY's avatar
PanZezhongQY committed
5
6
    infiniopHandle_t,
    infiniopTensorDescriptor_t,
xgqdut2016's avatar
xgqdut2016 committed
7
8
9
    open_lib,
    to_tensor,
    get_test_devices,
PanZezhongQY's avatar
PanZezhongQY committed
10
    check_error,
xgqdut2016's avatar
xgqdut2016 committed
11
12
13
14
15
16
    rearrange_if_needed,
    test_operator,
    get_args,
    debug,
    get_tolerance,
    profile_operation,
17
    create_workspace,
PanZezhongQY's avatar
PanZezhongQY committed
18
)
xgqdut2016's avatar
xgqdut2016 committed
19
from enum import Enum, auto
PanZezhongQY's avatar
PanZezhongQY committed
20

xgqdut2016's avatar
xgqdut2016 committed
21
22
23
24
# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
25
_TEST_CASES_ = [
26
    # shape, a_stride, b_stride, c_stride
27
28
    ((13, 4), None, None, None),
    ((13, 4), (10, 1), (10, 1), (10, 1)),
29
    ((13, 4), (0, 1), None, None),
30
31
    ((13, 4, 4), None, None, None),
    ((13, 4, 4), (20, 4, 1), (20, 4, 1), (20, 4, 1)),
32
    ((13, 4, 4), (4, 0, 1), (0, 4, 1), None),
33
34
    ((16, 5632), None, None, None),
    ((16, 5632), (13312, 1), (13312, 1), (13312, 1)),
35
36
    ((4, 4, 5632), None, None, None),
    ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), (45056, 5632, 1)),
37
38
]

39

40
41
42
43
44
45
class Inplace(Enum):
    OUT_OF_PLACE = auto()
    INPLACE_A = auto()
    INPLACE_B = auto()


46
47
# Inplace options applied for each test case in _TEST_CASES_
_INPLACE = [
48
49
50
    Inplace.OUT_OF_PLACE,
    Inplace.INPLACE_A,
    Inplace.INPLACE_B,
51
52
53
]

# Form the test cases by appending each element of _INPLACE to each tuple in _TEST_CASES_
xgqdut2016's avatar
xgqdut2016 committed
54
_TEST_CASES = [
55
56
57
    test_case + (inplace_item,)
    for test_case in _TEST_CASES_
    for inplace_item in _INPLACE
xgqdut2016's avatar
xgqdut2016 committed
58
]
xgqdut2016's avatar
xgqdut2016 committed
59

xgqdut2016's avatar
xgqdut2016 committed
60
# Data types used for testing
61
_TENSOR_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
xgqdut2016's avatar
xgqdut2016 committed
62
63
64

# Tolerance map for different data types
_TOLERANCE_MAP = {
65
    torch.float16: {"atol": 1e-3, "rtol": 1e-3},
66
    torch.bfloat16: {"atol": 5e-3, "rtol": 5e-3},
67
    torch.float32: {"atol": 2e-7, "rtol": 1e-7},
xgqdut2016's avatar
xgqdut2016 committed
68
69
70
71
72
73
}

DEBUG = False
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000
PanZezhongQY's avatar
PanZezhongQY committed
74

xgqdut2016's avatar
xgqdut2016 committed
75

PanZezhongQY's avatar
PanZezhongQY committed
76
77
78
79
80
81
82
83
84
class SwiGLUDescriptor(Structure):
    _fields_ = [("device", c_int32)]


infiniopSwiGLUDescriptor_t = POINTER(SwiGLUDescriptor)


def swiglu(a, b):
    return a * b / (1 + torch.exp(-b.float()).to(b.dtype))
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115


def process_tensors(c, c_strides, a, a_stride, b, b_stride, inplace):
    """
    rearrange the tensors if needed and apply the inplace config.
    if inplace is true and the output (i.e., c) is placed to the broadcasted input,
    the inplace config is ignored and out-of-place is used
    """
    original_c_strides = c_strides if c_strides else c.stride()

    def _rearrange(tensor, strides):
        if strides and 0 in strides:
            tensor.set_(tensor.untyped_storage(), 0, tensor.shape, strides)
            return tensor
        else:
            return rearrange_if_needed(tensor, strides)

    a, b, c = [
        _rearrange(tensor, stride)
        for tensor, stride in zip([a, b, c], [a_stride, b_stride, c_strides])
    ]
    c = (
        c
        if inplace == Inplace.OUT_OF_PLACE
        else (a if inplace == Inplace.INPLACE_A else b)
    )
    # if inplace is true and c has broadcasted config, reset it to the original unbroadcasted strides
    if 0 in c.stride():
        c.set_(c.untyped_storage(), 0, c.shape, original_c_strides)

    return a, b, c
PanZezhongQY's avatar
PanZezhongQY committed
116

117

xgqdut2016's avatar
xgqdut2016 committed
118
def test(
PanZezhongQY's avatar
PanZezhongQY committed
119
120
121
122
123
124
125
    lib,
    handle,
    torch_device,
    shape,
    a_stride=None,
    b_stride=None,
    c_stride=None,
xgqdut2016's avatar
xgqdut2016 committed
126
    inplace=Inplace.OUT_OF_PLACE,
PanZezhongQY's avatar
PanZezhongQY committed
127
128
129
130
    dtype=torch.float16,
    sync=None,
):
    print(
131
132
        f"Testing SwiGLU on {torch_device} with shape:{shape} a_stride:{a_stride} b_stride:{b_stride} c_stride:{c_stride} "
        f"dtype:{dtype} inplace:{inplace}"
PanZezhongQY's avatar
PanZezhongQY committed
133
    )
xgqdut2016's avatar
xgqdut2016 committed
134

PanZezhongQY's avatar
PanZezhongQY committed
135
136
    a = torch.rand(shape, dtype=dtype).to(torch_device)
    b = torch.rand(shape, dtype=dtype).to(torch_device)
137
    c = torch.rand(shape, dtype=dtype).to(torch_device)
138
    a, b, c = process_tensors(c, c_stride, a, a_stride, b, b_stride, inplace)
PanZezhongQY's avatar
PanZezhongQY committed
139
140
141

    ans = swiglu(a, b)

xgqdut2016's avatar
xgqdut2016 committed
142
143
144
145
146
147
    a_tensor, b_tensor = [to_tensor(tensor, lib) for tensor in [a, b]]
    c_tensor = (
        to_tensor(c, lib)
        if inplace == Inplace.OUT_OF_PLACE
        else (a_tensor if inplace == Inplace.INPLACE_A else b_tensor)
    )
PanZezhongQY's avatar
PanZezhongQY committed
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
    if sync is not None:
        sync()

    descriptor = infiniopSwiGLUDescriptor_t()
    check_error(
        lib.infiniopCreateSwiGLUDescriptor(
            handle,
            ctypes.byref(descriptor),
            c_tensor.descriptor,
            a_tensor.descriptor,
            b_tensor.descriptor,
        )
    )

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
xgqdut2016's avatar
xgqdut2016 committed
163
    for tensor in [a_tensor, b_tensor, c_tensor]:
164
        tensor.destroyDesc(lib)
xgqdut2016's avatar
xgqdut2016 committed
165

166
167
168
169
170
171
    workspace_size = c_uint64(0)
    check_error(
        lib.infiniopGetSwiGLUWorkspaceSize(descriptor, ctypes.byref(workspace_size))
    )
    workspace = create_workspace(workspace_size.value, c.device)

xgqdut2016's avatar
xgqdut2016 committed
172
173
174
    def lib_swiglu():
        check_error(
            lib.infiniopSwiGLU(
175
                descriptor,
176
177
                workspace.data_ptr() if workspace is not None else None,
                workspace_size.value,
178
179
180
181
                c_tensor.data,
                a_tensor.data,
                b_tensor.data,
                None,
xgqdut2016's avatar
xgqdut2016 committed
182
            )
PanZezhongQY's avatar
PanZezhongQY committed
183
        )
xgqdut2016's avatar
xgqdut2016 committed
184

xgqdut2016's avatar
xgqdut2016 committed
185
    lib_swiglu()
PanZezhongQY's avatar
PanZezhongQY committed
186

xgqdut2016's avatar
xgqdut2016 committed
187
188
189
190
    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
    if DEBUG:
        debug(c, ans, atol=atol, rtol=rtol)
    assert torch.allclose(c, ans, atol=atol, rtol=rtol)
PanZezhongQY's avatar
PanZezhongQY committed
191

xgqdut2016's avatar
xgqdut2016 committed
192
193
194
195
196
197
    # Profiling workflow
    if PROFILE:
        # fmt: off
        profile_operation("PyTorch", lambda: swiglu(a, b), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        profile_operation("    lib", lambda: lib_swiglu(), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        # fmt: on
PanZezhongQY's avatar
PanZezhongQY committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    check_error(lib.infiniopDestroySwiGLUDescriptor(descriptor))


if __name__ == "__main__":
    args = get_args()
    lib = open_lib()

    lib.infiniopCreateSwiGLUDescriptor.restype = c_int32
    lib.infiniopCreateSwiGLUDescriptor.argtypes = [
        infiniopHandle_t,
        POINTER(infiniopSwiGLUDescriptor_t),
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
    ]

214
215
216
217
218
219
    lib.infiniopGetSwiGLUWorkspaceSize.restype = c_int32
    lib.infiniopGetSwiGLUWorkspaceSize.argtypes = [
        infiniopSwiGLUDescriptor_t,
        POINTER(c_uint64),
    ]

PanZezhongQY's avatar
PanZezhongQY committed
220
221
222
223
    lib.infiniopSwiGLU.restype = c_int32
    lib.infiniopSwiGLU.argtypes = [
        infiniopSwiGLUDescriptor_t,
        c_void_p,
224
225
        c_uint64,
        c_void_p,
PanZezhongQY's avatar
PanZezhongQY committed
226
227
228
229
230
231
232
233
234
        c_void_p,
        c_void_p,
        c_void_p,
    ]

    lib.infiniopDestroySwiGLUDescriptor.restype = c_int32
    lib.infiniopDestroySwiGLUDescriptor.argtypes = [
        infiniopSwiGLUDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
235

xgqdut2016's avatar
xgqdut2016 committed
236
237
238
239
240
    # Configure testing options
    DEBUG = args.debug
    PROFILE = args.profile
    NUM_PRERUN = args.num_prerun
    NUM_ITERATIONS = args.num_iterations
xgqdut2016's avatar
xgqdut2016 committed
241

xgqdut2016's avatar
xgqdut2016 committed
242
243
    for device in get_test_devices(args):
        test_operator(lib, device, test, _TEST_CASES, _TENSOR_DTYPES)
PanZezhongQY's avatar
PanZezhongQY committed
244
245

    print("\033[92mTest passed!\033[0m")