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

# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
PanZezhong's avatar
PanZezhong committed
26
27
28
_TEST_CASES_ = [
    # (shape, x_strides, y_strides)
    ((1, 32, 128), None, None),
PanZezhong's avatar
PanZezhong committed
29
    ((10, 32, 64), None, None),
xgqdut2016's avatar
xgqdut2016 committed
30
31
    # 昇腾暂不满足这个用例,最后一维度 <=32 会有问题,可能与其核心
    # 接口 GatherMask 的内部实现相关,目前 48 64 128 都可以支持
PanZezhong's avatar
PanZezhong committed
32
33
34
    ((4, 1, 32), (64, 64, 1), None),
    ((11, 33, 128), None, (8000, 200, 1)),
    ((3, 32, 128), (8000, 200, 1), (7000, 128, 1)),
xgqdut2016's avatar
xgqdut2016 committed
35
36
37
]

# Data types used for testing
PanZezhong's avatar
PanZezhong committed
38
_TENSOR_DTYPES = [torch.float16, torch.float32]
xgqdut2016's avatar
xgqdut2016 committed
39
40
41

# Tolerance map for different data types
_TOLERANCE_MAP = {
PanZezhong's avatar
PanZezhong committed
42
    torch.float16: {"atol": 1e-3, "rtol": 1e-2},
PanZezhong's avatar
PanZezhong committed
43
    torch.float32: {"atol": 1e-4, "rtol": 1e-3},
xgqdut2016's avatar
xgqdut2016 committed
44
}
PanZezhongQY's avatar
PanZezhongQY committed
45

PanZezhong's avatar
PanZezhong committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

class Inplace(Enum):
    OUT_OF_PLACE = auto()
    INPLACE_X = auto()


_INPLACE = [
    Inplace.OUT_OF_PLACE,
    Inplace.INPLACE_X,
]

_TEST_CASES = [
    test_case + (inplace_item,)
    for test_case in _TEST_CASES_
    for inplace_item in _INPLACE
]

63
64
65
66
67
DEBUG = False
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000

PanZezhongQY's avatar
PanZezhongQY committed
68
69
70
71
72
73
74
75

class RoPEDescriptor(Structure):
    _fields_ = [("device", c_int32)]


infiniopRoPEDescriptor_t = POINTER(RoPEDescriptor)


PanZezhong's avatar
PanZezhong committed
76
def rotary_embedding(t, sin, cos, torch_device):
PanZezhongQY's avatar
PanZezhongQY committed
77
    dh = t.shape[2]
PanZezhong's avatar
PanZezhong committed
78
    dt = t.dtype
79
    assert dh % 2 == 0, "Embedding dimension must be even."
PanZezhong's avatar
PanZezhong committed
80
81
82
83
84
85
86
87
88
89
90
    t_even = t[..., 0::2]  # [seq_len, n_head, dh // 2]
    t_odd = t[..., 1::2]  # [seq_len, n_head, dh // 2]
    cos = cos.unsqueeze(1)  # [seq_len, 1, dh // 2]
    sin = sin.unsqueeze(1)  # [seq_len, 1, dh // 2]
    if torch_device == "cpu":
        (t_even, t_odd, cos, sin) = (
            t_even.float(),
            t_odd.float(),
            cos.float(),
            sin.float(),
        )
91
92
93
94
95
96
97

    t_out_even = t_even * cos - t_odd * sin
    t_out_odd = t_even * sin + t_odd * cos

    t_out = torch.empty_like(t)
    t_out[..., 0::2] = t_out_even
    t_out[..., 1::2] = t_out_odd
98

PanZezhong's avatar
PanZezhong committed
99
    return t_out.to(dt).to(torch_device)
PanZezhongQY's avatar
PanZezhongQY committed
100

101

PanZezhong's avatar
PanZezhong committed
102
103
def sin_cos_table(pos, dim, torch_device, theta, dtype):
    assert dim % 2 == 0, "Embedding dimension must be even."
PanZezhongQY's avatar
PanZezhongQY committed
104
105
106
107
    freqs = (1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))).to(
        torch_device
    )
    angles = torch.outer(pos, freqs)
PanZezhong's avatar
PanZezhong committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    return torch.sin(angles).to(dtype), torch.cos(angles).to(dtype)


def test(
    lib,
    handle,
    torch_device,
    shape,
    x_strides=None,
    y_strides=None,
    inplace=Inplace.OUT_OF_PLACE,
    dtype=torch.float32,
):
    if inplace == Inplace.INPLACE_X:
        y_strides = x_strides
PanZezhongQY's avatar
PanZezhongQY committed
123
    print(
PanZezhong's avatar
PanZezhong committed
124
        f"Testing Rotary Positional Embedding on {torch_device} with shape:{shape} x_strides:{x_strides} y_strides:{y_strides} and dtype:{dtype} inplace:{inplace}"
PanZezhongQY's avatar
PanZezhongQY committed
125
126
    )

PanZezhong's avatar
PanZezhong committed
127
128
129
130
131
132
133
134
135
136
    x = torch.rand(shape, dtype=dtype).to(torch_device)
    x = rearrange_if_needed(x, x_strides)
    if inplace == Inplace.INPLACE_X:
        y = x
    else:
        y = torch.rand(shape, dtype=dtype).to(torch_device)
        y = rearrange_if_needed(y, y_strides)
    theta = 1e5
    pos = torch.arange(0, x.shape[0], dtype=torch.int32).to(torch_device)
    sin_table, cos_table = sin_cos_table(pos, x.shape[2], x.device, theta, dtype)
xgqdut2016's avatar
xgqdut2016 committed
137

PanZezhong's avatar
PanZezhong committed
138
    ans = rotary_embedding(x, sin_table, cos_table, torch_device)
PanZezhongQY's avatar
PanZezhongQY committed
139
140

    descriptor = infiniopRoPEDescriptor_t()
PanZezhong's avatar
PanZezhong committed
141
142
143
    x_tensor, pos_tensor, sin_table_tensor, cos_table_tensor = [
        to_tensor(tensor, lib, force_unsigned=True)
        for tensor in [x, pos, sin_table, cos_table]
xgqdut2016's avatar
xgqdut2016 committed
144
    ]
PanZezhong's avatar
PanZezhong committed
145
146
147
148
    if inplace == Inplace.INPLACE_X:
        y_tensor = x_tensor
    else:
        y_tensor = to_tensor(y, lib)
PanZezhongQY's avatar
PanZezhongQY committed
149
150

    if torch_device == "npu":
xgqdut2016's avatar
xgqdut2016 committed
151
        synchronize_device(torch_device)
PanZezhongQY's avatar
PanZezhongQY committed
152
153
154
155

    check_error(
        lib.infiniopCreateRoPEDescriptor(
            handle,
156
            ctypes.byref(descriptor),
PanZezhong's avatar
PanZezhong committed
157
158
            y_tensor.descriptor,
            x_tensor.descriptor,
PanZezhongQY's avatar
PanZezhongQY committed
159
160
161
162
163
164
165
            pos_tensor.descriptor,
            sin_table_tensor.descriptor,
            cos_table_tensor.descriptor,
        )
    )

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
PanZezhong's avatar
PanZezhong committed
166
167
    for tensor in [y_tensor, x_tensor, pos_tensor, sin_table_tensor, cos_table_tensor]:
        tensor.destroyDesc(lib)
PanZezhongQY's avatar
PanZezhongQY committed
168
169
170
171
172

    workspace_size = c_uint64(0)
    check_error(
        lib.infiniopGetRoPEWorkspaceSize(descriptor, ctypes.byref(workspace_size))
    )
PanZezhong's avatar
PanZezhong committed
173
    workspace = create_workspace(workspace_size.value, x.device)
174
175
176
177
178
179
180

    def lib_rope():
        check_error(
            lib.infiniopRoPE(
                descriptor,
                workspace.data_ptr() if workspace is not None else None,
                workspace_size.value,
PanZezhong's avatar
PanZezhong committed
181
182
                y_tensor.data,
                x_tensor.data,
183
184
185
186
187
                pos_tensor.data,
                sin_table_tensor.data,
                cos_table_tensor.data,
                None,
            )
PanZezhongQY's avatar
PanZezhongQY committed
188
189
        )

190
    lib_rope()
xgqdut2016's avatar
xgqdut2016 committed
191

xgqdut2016's avatar
xgqdut2016 committed
192
    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
193
    if DEBUG:
PanZezhong's avatar
PanZezhong committed
194
195
        debug(y, ans, atol=atol, rtol=rtol)
    assert torch.allclose(y, ans, atol=atol, rtol=rtol)
xgqdut2016's avatar
xgqdut2016 committed
196

197
198
199
    if PROFILE:
        profile_operation(
            "PyTorch",
PanZezhong's avatar
PanZezhong committed
200
            lambda: rotary_embedding(x, pos, theta, torch_device),
201
202
203
204
205
206
207
            torch_device,
            NUM_PRERUN,
            NUM_ITERATIONS,
        )
        profile_operation(
            "    lib", lambda: lib_rope(), torch_device, NUM_PRERUN, NUM_ITERATIONS
        )
PanZezhongQY's avatar
PanZezhongQY committed
208

209
    check_error(lib.infiniopDestroyRoPEDescriptor(descriptor))
PanZezhongQY's avatar
PanZezhongQY committed
210

211

PanZezhongQY's avatar
PanZezhongQY committed
212
213
214
if __name__ == "__main__":
    args = get_args()
    lib = open_lib()
xgqdut2016's avatar
xgqdut2016 committed
215

PanZezhongQY's avatar
PanZezhongQY committed
216
217
218
219
220
221
222
223
224
    lib.infiniopCreateRoPEDescriptor.restype = c_int32
    lib.infiniopCreateRoPEDescriptor.argtypes = [
        infiniopHandle_t,
        POINTER(infiniopRoPEDescriptor_t),
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
225

PanZezhongQY's avatar
PanZezhongQY committed
226
227
228
229
230
    lib.infiniopGetRoPEWorkspaceSize.restype = c_int32
    lib.infiniopGetRoPEWorkspaceSize.argtypes = [
        infiniopRoPEDescriptor_t,
        POINTER(c_uint64),
    ]
xgqdut2016's avatar
xgqdut2016 committed
231

PanZezhongQY's avatar
PanZezhongQY committed
232
233
234
235
236
237
238
239
240
241
242
    lib.infiniopRoPE.restype = c_int32
    lib.infiniopRoPE.argtypes = [
        infiniopRoPEDescriptor_t,
        c_void_p,
        c_uint64,
        c_void_p,
        c_void_p,
        c_void_p,
        c_void_p,
        c_void_p,
    ]
xgqdut2016's avatar
xgqdut2016 committed
243

PanZezhongQY's avatar
PanZezhongQY committed
244
245
246
247
    lib.infiniopDestroyRoPEDescriptor.restype = c_int32
    lib.infiniopDestroyRoPEDescriptor.argtypes = [
        infiniopRoPEDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
248

249
250
251
252
253
254
255
256
    # Configure testing options
    DEBUG = args.debug
    PROFILE = args.profile
    NUM_PRERUN = args.num_prerun
    NUM_ITERATIONS = args.num_iterations

    # Execute tests
    for device in get_test_devices(args):
xgqdut2016's avatar
xgqdut2016 committed
257
        test_operator(lib, device, test, _TEST_CASES, _TENSOR_DTYPES)
258

PanZezhongQY's avatar
PanZezhongQY committed
259
    print("\033[92mTest passed!\033[0m")