CodeAlign-Runtime/benchmarks/gemv_benchmark.cpp at main · Sergasgr/CodeAlign-Runtime · GitHub

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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#include "../src/gemv/gemv.h"
#include "benchmark_utils.h"
#include <cuda_runtime.h>
#include <iostream>
#include <vector>
#include <cstdint>
#include <cmath>

int main() {
    int rows = 4864, cols = 896; // MLP Up projection dimensions for Qwen2.5-0.5B

    size_t bytes_mat = rows * cols * sizeof(float);
    size_t bytes_vec = cols * sizeof(float);
    size_t bytes_out = rows * sizeof(float);

    std::vector<float> h_mat(rows * cols, 1.0f);
    std::vector<float> h_vec(cols, 1.0f);
    std::vector<float> h_out_naive(rows, 1.0f);
    std::vector<float> h_out_opt(rows, 1.0f);
    std::vector<uint32_t> h_q_mat;
    std::vector<float> h_scales;
    std::vector<float> h_out_q(rows, 1.0f);

    symmetric_quantization(h_mat, h_q_mat, h_scales, rows, cols);

    size_t bytes_q_mat = h_q_mat.size() * sizeof(uint32_t);
    size_t bytes_scales = h_scales.size() * sizeof(float);

    float *d_mat, *d_vec, *d_out_naive, *d_out_opt, *d_scales, *d_out_q;
    uint32_t *d_q_mat;

    cudaMalloc((void**)&d_mat, bytes_mat);
    cudaMalloc((void**)&d_vec, bytes_vec);
    cudaMalloc((void**)&d_out_naive, bytes_out);
    cudaMalloc((void**)&d_out_opt, bytes_out);
    cudaMalloc((void**)&d_q_mat, bytes_q_mat);
    cudaMalloc((void**)&d_scales, bytes_scales);
    cudaMalloc((void**)&d_out_q, bytes_out);

    cudaMemcpy(d_mat, h_mat.data(), bytes_mat, cudaMemcpyHostToDevice);
    cudaMemcpy(d_vec, h_vec.data(), bytes_vec, cudaMemcpyHostToDevice);
    cudaMemcpy(d_out_naive, h_out_naive.data(), bytes_out, cudaMemcpyHostToDevice);
    cudaMemcpy(d_out_opt, h_out_opt.data(), bytes_out, cudaMemcpyHostToDevice);
    cudaMemcpy(d_q_mat, h_q_mat.data(), bytes_q_mat, cudaMemcpyHostToDevice);
    cudaMemcpy(d_scales, h_scales.data(), bytes_scales, cudaMemcpyHostToDevice);
    cudaMemcpy(d_out_q, h_out_q.data(), bytes_out, cudaMemcpyHostToDevice);

    benchmark_kernel("NAIVE KERNEL", [&]() {
        run_gemv_naive(d_mat, d_vec, d_out_naive, rows, cols);
    }, bytes_mat + bytes_vec + bytes_out);

    cudaMemcpy(h_out_naive.data(), d_out_naive, bytes_out, cudaMemcpyDeviceToHost);

    benchmark_kernel("OPTIMIZED KERNEL", [&]() {
        run_gemv_optimized(d_mat, d_vec, d_out_opt, rows, cols);
    }, bytes_mat + bytes_vec + bytes_out);

    cudaMemcpy(h_out_opt.data(), d_out_opt, bytes_out, cudaMemcpyDeviceToHost);

    benchmark_kernel("QUANTIZED NAIVE KERNEL", [&]() {
        run_gemv_int4_naive_kernel(d_q_mat, d_scales, d_vec, d_out_q, rows, cols);
    }, bytes_q_mat + bytes_scales + bytes_vec + bytes_out);

    cudaMemcpy(h_out_q.data(), d_out_q, bytes_out, cudaMemcpyDeviceToHost);

    benchmark_kernel("QUANTIZED OPTIMIZED KERNEL", [&]() {
        run_gemv_int4_optimized_kernel(d_q_mat, d_scales, d_vec, d_out_q, rows, cols);
    }, bytes_q_mat + bytes_scales + bytes_vec + bytes_out);

    cudaMemcpy(h_out_q.data(), d_out_q, bytes_out, cudaMemcpyDeviceToHost);

    for(int i = 0; i < rows; i++) {
        if (std::abs(h_out_naive[i] - h_out_opt[i]) > 1e-4 || std::abs(h_out_naive[i] - h_out_q[i]) > 0.05f)  {
            std::cout << "Validation Error in index " << i << "\n";
            std::cout << "Naive: " << h_out_naive[i]
                      << " Opt: " << h_out_opt[i]
                      << " Q: " << h_out_q[i] << "\n";
            break;
        }
    }

    cudaFree(d_mat);
    cudaFree(d_vec);
    cudaFree(d_out_naive);
    cudaFree(d_out_opt);
    cudaFree(d_q_mat);
    cudaFree(d_scales);
    cudaFree(d_out_q);

    return 0;
}