File size: 2,322 Bytes
4a28d4d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72


#include "src/turbomind/core/allocator.h"
#include "src/turbomind/core/context.h"
#include "src/turbomind/core/data_type.h"
#include "src/turbomind/core/stream.h"

#include "src/turbomind/kernels/gemm/test/test_utils.h"
#include "src/turbomind/kernels/quantization.h"

using namespace turbomind;

int main()
{
    core::ContextGuard ctx{core::Stream::create(), core::Allocator{kCPU}, core::Allocator{kDEVICE}};

    auto stream = core::Context::stream().handle();

    const int m = 1024, n = 2048, gs = 128;

    Tensor_<bfloat16_t> h_x{{m, n}, kCPU};
    Tensor_<bfloat16_t> h_x_f{{m, n}, kCPU};

    Tensor_<bfloat16_t> x{{m, n}, kDEVICE};
    Tensor_<bfloat16_t> x_f{{m, n}, kDEVICE};
    Tensor_<fp8_e4m3_t> x_q{{m, n}, kDEVICE};

    // Tensor_<float> x_s{{{m, n / gs}, {1, round_up(m, 4)}}, kDEVICE};
    Tensor_<float> x_s;

    RNG r;
    r.set_stream(stream);

    /////////////////////////////////////////////////////////////////////////////////////
    // round trip of dequant(quant(x))
    r.UniformFloat(x, 2.f, 2.f);  // [-1, +1]
    Copy(x, h_x);
    QuantizeSymm(x_q, x_s, x, stream);
    DequantizeSymm(x_f, x_q, x_s, stream);
    Copy(x_f, h_x_f);
    FC_Header();
    FC_Print(FastCompare(x_f, x, stream));

    /////////////////////////////////////////////////////////////////////////////////////
    // round trip of dequant(quant(dequant(quant(x)))), aligned representable values
    Copy(x_f, x);
    Clear(x_f);
    QuantizeSymm(x_q, x_s, x, stream);
    DequantizeSymm(x_f, x_q, x_s, stream);
    FC_Print(FastCompare(x_f, x, stream));

    /////////////////////////////////////////////////////////////////////////////////////
    // round trip of dequant(quant(x))
    // x_s = {{cdiv(m, gs), cdiv(n, gs)}, kDEVICE};
    x_s = {};
    r.UniformFloat(x, 2.f, 2.f);  // [-1, +1]
    Copy(x, h_x);
    QuantizeSymmBlock(x_q, x_s, x, stream);
    DequantizeSymmBlock(x_f, x_q, x_s, stream);
    FC_Print(FastCompare(x_f, x, stream));

    /////////////////////////////////////////////////////////////////////////////////////
    // round trip of dequant(quant(dequant(quant(x)))), aligned representable values
    Copy(x_f, x);
    Clear(x_f);
    QuantizeSymmBlock(x_q, x_s, x, stream);
    DequantizeSymmBlock(x_f, x_q, x_s, stream);
    FC_Print(FastCompare(x_f, x, stream));

    return 0;
}