Bash-Math-LLM/full_kernels.awk at main · nareshis21/Bash-Math-LLM · 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
92
93
94
95
96
# full_kernels.awk - Comprehensive Qwen3 Inference kernels

function decode_fp16(b1, b2) {
    if (b1 == "" || b2 == "") return 0
    local_sign = (b2 >= 128) ? -1 : 1
    local_exp = int((b2 % 128) / 4)
    local_mant = (b2 % 4) * 256 + b1
    if (local_exp == 0) return local_sign * (local_mant / 1024) * (2^-14)
    if (local_exp == 31) return (local_mant == 0) ? (local_sign * 1e38) : 0
    return local_sign * (1 + local_mant / 1024) * (2^(local_exp - 15))
}

# Mode: decode
# Turns stdin binary (od) into lines of floats
function decode_bin() {
    for (i=1; i<=NF; i++) {
        buf[buf_ptr++] = $i
        if (buf_ptr == 2) {
            val = decode_fp16(buf[0], buf[1])
            print val
            buf_ptr = 0
        }
    }
}

function matmul() {
    if (init == 0) {
        idx = 0; while ((getline line < v_file) > 0) v_in[idx++] = line; close(v_file)
        for (i=0; i<out_dim; i++) v_out[i] = 0
        init = 1; curr_row = 0; curr_col = 0; buf_ptr = 0; count = 0
    }
    for (i=1; i<=NF; i++) {
        buf[buf_ptr++] = $i
        if (buf_ptr == 2) {
            val = decode_fp16(buf[0], buf[1])
            res = val * v_in[curr_col]
            v_out[curr_row] += res

            if (visual && count % 1000 == 0) {
                total_ops = in_dim * out_dim
                pct = (count / total_ops) * 100
                bar_len = 15; filled = int((count / total_ops) * bar_len)
                bar = ""; for (b=0; b<bar_len; b++) bar = bar (b < filled ? "█" : "░")
                C_CYAN = "\033[36m"; C_GREEN = "\033[32m"; C_RESET = "\033[0m"; C_DIM = "\033[2m"; C_BLUE = "\033[34m"
                # Print a single-line progress bar with Layer info
                l_info = (layer != "" ? "L" layer "/" max_layers " " : "")
                printf " %s%s[%s]%s %2d%% | %s%7.4f %s× %s%7.4f %s→ %s%10.6f %s(R%d)         \r", \
                    C_CYAN, l_info, bar, C_RESET, pct, C_GREEN, val, C_DIM, C_BLUE, v_in[curr_col], C_DIM, C_GREEN, res, C_RESET, curr_row > "/dev/stderr"
            }
            count++; curr_col++
            if (curr_col == in_dim) { curr_col = 0; curr_row++ }
            buf_ptr = 0
        }
    }
}

BEGIN {
    if (mode == "rmsnorm") {
        idx = 0; while ((getline line < v_file) > 0) v[idx++] = line; close(v_file)
        idx = 0; while ((getline line < w_file) > 0) w[idx++] = line; close(w_file)
        ss = 0; for (i=0; i<dim; i++) ss += v[i] * v[i]
        inv_rms = 1.0 / sqrt(ss / dim + eps)
        for (i=0; i<dim; i++) print (v[i] * inv_rms) * w[i]
        exit
    }
    if (mode == "silu") {
        idx = 0; while ((getline line < gate_file) > 0) g[idx++] = line; close(gate_file)
        idx = 0; while ((getline line < up_file) > 0) u[idx++] = line; close(up_file)
        for (i=0; i<dim; i++) {
            sig = 1.0 / (1.0 + exp(-g[i]))
            print (g[i] * sig) * u[i]
        }
        exit
    }
    if (mode == "argmax") {
        idx = 0; max_val = -1e38; winner = 0
        while ((getline line < f1) > 0) {
            if (line > max_val) { max_val = line; winner = idx }
            idx++
        }
        close(f1)
        print winner
        exit
    }
}

{
    if (mode == "matmul") matmul()
    if (mode == "decode") decode_bin()
}

END {
    if (mode == "matmul") {
        for (i=0; i<out_dim; i++) print v_out[i]
    }
}