test_bench3_bias.py 2.89 KB
Newer Older
helloyongyang's avatar
helloyongyang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
from lightx2v_kernel.gemm import scaled_fp4_quant, cutlass_scaled_fp4_mm
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")
        bias = torch.randn(1, n, dtype=torch.bfloat16).cuda()

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

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

helloyongyang's avatar
helloyongyang committed
18
19
20
21
22
23
        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
24

helloyongyang's avatar
helloyongyang committed
25
26
27
28
29
30
        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
31

helloyongyang's avatar
helloyongyang committed
32
33
34
35
36
37
38
39
40
41
        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
42
            ref_output_tensor = linear(input_tensor)
helloyongyang's avatar
helloyongyang committed
43
44
        torch.cuda.synchronize()
        end_time = time.time()
helloyongyang's avatar
fix ci  
helloyongyang committed
45

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

helloyongyang's avatar
helloyongyang committed
49
50
51
52
53
54
55
56
        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")
        bias = torch.randn(1, n, dtype=torch.bfloat16).cuda()
helloyongyang's avatar
fix ci  
helloyongyang committed
57

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

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

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

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

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

helloyongyang's avatar
helloyongyang committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
        # 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
90

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