random_sample.py 5.29 KB
Newer Older
xgqdut2016's avatar
xgqdut2016 committed
1
import torch
PanZezhongQY's avatar
PanZezhongQY committed
2
import ctypes
3
from ctypes import c_uint64
xgqdut2016's avatar
xgqdut2016 committed
4
from libinfiniop import (
5
6
    LIBINFINIOP,
    TestTensor,
xgqdut2016's avatar
xgqdut2016 committed
7
    get_test_devices,
PanZezhongQY's avatar
PanZezhongQY committed
8
    check_error,
xgqdut2016's avatar
xgqdut2016 committed
9
10
    test_operator,
    get_args,
xgqdut2016's avatar
xgqdut2016 committed
11
    debug_all,
xgqdut2016's avatar
xgqdut2016 committed
12
13
    get_tolerance,
    profile_operation,
14
15
16
17
18
    TestWorkspace,
    InfiniDtype,
    InfiniDtypeNames,
    InfiniDeviceNames,
    infiniopOperatorDescriptor_t,
PanZezhongQY's avatar
PanZezhongQY committed
19
20
)

xgqdut2016's avatar
xgqdut2016 committed
21
22
23
24
25
26
# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
_TEST_CASES = [
    # voc, random_val, topp, topk, temperature
xgqdut2016's avatar
xgqdut2016 committed
27
28
29
30
31
32
33
34
35
36
    (512, 0.8, 0.8, 3, 0.5),
    (4096, 0.05, 0.9, 5, 1.0),
    (16384, 0.15, 0.85, 10, 2.0),
    (512, 0.08, 0, 3, 0.5),
    (4096, 0.5, 0.9, 1, 1.0),
    (16384, 0.15, 0, 1, 2.0),
    (16384, 0.15, 0, 1, 2.0),
    (32000, 0.08, 0.8, 50, 1.0),
    (32000, 0.08, 1.0, 25, 1.0),
    # (119696, 0.01, 1.0, 100, 1.0),
xgqdut2016's avatar
xgqdut2016 committed
37
38
39
]

# Data types used for testing
40
_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16]
xgqdut2016's avatar
xgqdut2016 committed
41
42

_TOLERANCE_MAP = {
43
44
    InfiniDtype.F16: {"atol": 0, "rtol": 0},
    InfiniDtype.BF16: {"atol": 0, "rtol": 0},
xgqdut2016's avatar
xgqdut2016 committed
45
}
xgqdut2016's avatar
xgqdut2016 committed
46

47

xgqdut2016's avatar
xgqdut2016 committed
48
DEBUG = False
xgqdut2016's avatar
xgqdut2016 committed
49
50
51
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000
PanZezhongQY's avatar
PanZezhongQY committed
52
53


54
def random_sample(data, random_val, topp, topk, voc, temperature):
55
    if topp > 0 and topk > 1:
56
        sorted_vals, sorted_indices = torch.sort(data, descending=True)
57

58
59
60
        scaled_vals = (sorted_vals - sorted_vals[0]) / temperature
        probs = torch.softmax(scaled_vals, dim=0)
        cum_probs = torch.cumsum(probs, dim=0)
61

62
63
        k_index = min(topk, voc) - 1
        threshold = min(cum_probs[k_index], topp) * random_val
64

65
66
67
68
69
        try:
            idx = torch.searchsorted(cum_probs, threshold)
        except Exception:
            # Fallback for manual search if torch.searchsorted is not supported
            indices = (cum_probs >= threshold).nonzero(as_tuple=True)[0]
70
71
72
73
74
            idx = (
                indices[0]
                if indices.numel() > 0
                else torch.tensor(len(cum_probs) - 1, device=cum_probs.device)
            )
75
        return sorted_indices[idx]
76

77
    return torch.argmax(data)
zhangyue's avatar
zhangyue committed
78
79


80
81
def test(
    handle,
82
    device,
83
84
85
86
87
    voc,
    random_val,
    topp,
    topk,
    temperature,
88
    dtype=InfiniDtype.F16,
89
    sync=None,
90
):
91
    print(
92
        f"Testing RandomSample on {InfiniDeviceNames[device]} with voc:{voc} random_val:{random_val} topp:{topp} topk:{topk} temperature:{temperature} dtype:{InfiniDtypeNames[dtype]}"
93
    )
xgqdut2016's avatar
xgqdut2016 committed
94

PanZezhongQY's avatar
PanZezhongQY committed
95
    _perm = torch.randperm(voc)
96
97
98
    logits = TestTensor.from_torch(
        torch.arange(voc)[_perm].float() * 0.0001, dtype, device
    )
99
100

    ans = random_sample(
101
        logits.torch_tensor(), random_val, topp, topk, voc, temperature
102
103
    ).to(
        torch.int32
104
    )  # 这个函数在device速度可能会很慢,可以通过data.to("cpu")方式加快计算过程
xgqdut2016's avatar
xgqdut2016 committed
105

106
    indices = TestTensor([], None, InfiniDtype.I32, device, mode="zeros")
PanZezhongQY's avatar
PanZezhongQY committed
107

108
109
110
    if sync is not None:
        sync()

111
    descriptor = infiniopOperatorDescriptor_t()
PanZezhongQY's avatar
PanZezhongQY committed
112
    check_error(
113
        LIBINFINIOP.infiniopCreateRandomSampleDescriptor(
114
115
            handle,
            ctypes.byref(descriptor),
116
117
            indices.descriptor,
            logits.descriptor,
PanZezhongQY's avatar
PanZezhongQY committed
118
119
120
121
        )
    )

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
122
123
    for tensor in [logits, indices]:
        tensor.destroy_desc()
PanZezhongQY's avatar
PanZezhongQY committed
124
125
126

    workspace_size = c_uint64(0)
    check_error(
127
        LIBINFINIOP.infiniopGetRandomSampleWorkspaceSize(
PanZezhongQY's avatar
PanZezhongQY committed
128
129
130
            descriptor, ctypes.byref(workspace_size)
        )
    )
131
    workspace = TestWorkspace(workspace_size.value, device)
132

xgqdut2016's avatar
xgqdut2016 committed
133
134
    def lib_random_sample():
        check_error(
135
            LIBINFINIOP.infiniopRandomSample(
xgqdut2016's avatar
xgqdut2016 committed
136
                descriptor,
137
                workspace.data(),
xgqdut2016's avatar
xgqdut2016 committed
138
                workspace_size.value,
139
140
                indices.data(),
                logits.data(),
xgqdut2016's avatar
xgqdut2016 committed
141
142
143
144
145
146
                random_val,
                topp,
                topk,
                temperature,
                None,
            )
PanZezhongQY's avatar
PanZezhongQY committed
147
148
        )

xgqdut2016's avatar
xgqdut2016 committed
149
150
    lib_random_sample()

151
152
    if sync is not None:
        sync()
xgqdut2016's avatar
xgqdut2016 committed
153
154
155
156

    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
    if DEBUG:
        debug_all(
157
158
            (indices.actual_tensor(), logits.actual_tensor()[indices.actual_tensor()]),
            (ans, logits.torch_tensor()[ans]),
xgqdut2016's avatar
xgqdut2016 committed
159
160
161
162
            "or",
            atol=atol,
            rtol=rtol,
        )
163
164
165
166
    assert (
        indices.actual_tensor() == ans
        or logits.actual_tensor()[indices.actual_tensor()] == logits.torch_tensor()[ans]
    )
xgqdut2016's avatar
xgqdut2016 committed
167

xgqdut2016's avatar
xgqdut2016 committed
168
169
170
    # Profiling workflow
    if PROFILE:
        # fmt: off
171
        profile_operation("PyTorch", lambda: random_sample(
172
173
174
            logits.torch_tensor(), random_val, topp, topk, voc, temperature
        ), device, NUM_PRERUN, NUM_ITERATIONS)
        profile_operation("    lib", lambda: lib_random_sample(), device, NUM_PRERUN, NUM_ITERATIONS)
xgqdut2016's avatar
xgqdut2016 committed
175
        # fmt: on
176
    check_error(LIBINFINIOP.infiniopDestroyRandomSampleDescriptor(descriptor))
PanZezhongQY's avatar
PanZezhongQY committed
177

178

PanZezhongQY's avatar
PanZezhongQY committed
179
180
181
if __name__ == "__main__":
    args = get_args()

xgqdut2016's avatar
xgqdut2016 committed
182
    DEBUG = args.debug
xgqdut2016's avatar
xgqdut2016 committed
183
184
185
186
187
188
    PROFILE = args.profile
    NUM_PRERUN = args.num_prerun
    NUM_ITERATIONS = args.num_iterations

    # Execute tests
    for device in get_test_devices(args):
189
        test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES)
190

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