test_bench3_bias.py 2.83 KB
Newer Older
helloyongyang's avatar
helloyongyang committed
1
2
3
4
5
6
7
8
9
import torch
import time
from test_bench2 import MMWeightFp4


def test_speed(m, k, n):
    with torch.no_grad():
        input_tensor = torch.randn(m, k, dtype=torch.bfloat16).cuda()
        weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
Xtra's avatar
Xtra committed
10
        bias = torch.ones(1, n, dtype=torch.bfloat16).cuda() * 50
helloyongyang's avatar
helloyongyang committed
11
12

        mm = MMWeightFp4(weight, bias)
helloyongyang's avatar
fix ci  
helloyongyang committed
13

helloyongyang's avatar
helloyongyang committed
14
15
        # warmup
        output_tensor = mm.apply(input_tensor)
helloyongyang's avatar
fix ci  
helloyongyang committed
16

helloyongyang's avatar
helloyongyang committed
17
18
19
20
21
22
        torch.cuda.synchronize()
        start_time = time.time()
        for i in range(100):
            output_tensor = mm.apply(input_tensor)
        torch.cuda.synchronize()
        end_time = time.time()
helloyongyang's avatar
fix ci  
helloyongyang committed
23

helloyongyang's avatar
helloyongyang committed
24
25
26
27
28
29
        lightx2v_kernel_time = (end_time - start_time) / 100
        print(f"lightx2v-kernel time: {lightx2v_kernel_time}")

        input_tensor = torch.randn(m, n, dtype=torch.bfloat16).cuda()
        weight = torch.randn(k, n, dtype=torch.bfloat16, device="cuda")
        bias = torch.randn(1, k, dtype=torch.bfloat16).cuda()
helloyongyang's avatar
fix ci  
helloyongyang committed
30

helloyongyang's avatar
helloyongyang committed
31
32
33
34
35
36
37
38
39
40
        linear = torch.nn.Linear(k, n, bias=True).cuda()
        linear.weight.data = weight
        linear.bias.data = bias

        # warmup
        ref_output_tensor = linear(input_tensor)

        torch.cuda.synchronize()
        start_time = time.time()
        for i in range(100):
helloyongyang's avatar
fix ci  
helloyongyang committed
41
            ref_output_tensor = linear(input_tensor)
helloyongyang's avatar
helloyongyang committed
42
43
        torch.cuda.synchronize()
        end_time = time.time()
helloyongyang's avatar
fix ci  
helloyongyang committed
44

helloyongyang's avatar
helloyongyang committed
45
46
        ref_time = (end_time - start_time) / 100
        print(f"ref time: {ref_time}")
helloyongyang's avatar
fix ci  
helloyongyang committed
47

helloyongyang's avatar
helloyongyang committed
48
49
50
51
52
53
54
        print(f"speedup: {ref_time / lightx2v_kernel_time:.3f}")


def test_accuracy(m, k, n):
    with torch.no_grad():
        input_tensor = torch.randn(m, k, dtype=torch.bfloat16).cuda()
        weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
Xtra's avatar
Xtra committed
55
        bias = torch.ones(1, n, dtype=torch.bfloat16).cuda() * 50
helloyongyang's avatar
fix ci  
helloyongyang committed
56

helloyongyang's avatar
helloyongyang committed
57
58
59
        linear = torch.nn.Linear(k, n, bias=True).cuda()
        linear.weight.data = weight
        linear.bias.data = bias
helloyongyang's avatar
fix ci  
helloyongyang committed
60

helloyongyang's avatar
helloyongyang committed
61
62
63
        ref_output_tensor = linear(input_tensor)

        mm = MMWeightFp4(weight, bias)
helloyongyang's avatar
fix ci  
helloyongyang committed
64

helloyongyang's avatar
helloyongyang committed
65
        output_tensor = mm.apply(input_tensor)
helloyongyang's avatar
fix ci  
helloyongyang committed
66

helloyongyang's avatar
helloyongyang committed
67
68
        # print(f"ref_output_tensor: {ref_output_tensor}")
        # print(f"output_tensor: {output_tensor}")
helloyongyang's avatar
fix ci  
helloyongyang committed
69

helloyongyang's avatar
helloyongyang committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
        # cosine
        cos = torch.nn.functional.cosine_similarity(ref_output_tensor.flatten(), output_tensor.flatten(), dim=0)
        print(f"cos : {cos}")


if __name__ == "__main__":
    test_sizes = [
        (32130, 5120, 5120),
        (512, 5120, 5120),
        (257, 5120, 5120),
        (32130, 5120, 13824),
        (32130, 13824, 5120),
        (75348, 5120, 5120),
        (75348, 13824, 5120),
        (32760, 1536, 1536),
        (512, 1536, 1536),
        (32760, 1536, 8960),
        (32760, 8960, 1536),
    ]
helloyongyang's avatar
fix ci  
helloyongyang committed
89

helloyongyang's avatar
helloyongyang committed
90
91
    for i, (m, k, n) in enumerate(test_sizes):
        print("-" * 30)
helloyongyang's avatar
fix ci  
helloyongyang committed
92
        print(f"测试 {i + 1}: 张量大小 ({m}, {k}, {n})")
helloyongyang's avatar
helloyongyang committed
93
94
        test_accuracy(m, k, n)
        test_speed(m, k, n)