`nx_qlayer.nx` source

↩ module page · 117 lines · 5101 B
1// nx_qlayer.nx -- the full decode LAYER: y = W . x, the operation an LLM does thousands of times
2// per token. Built honestly for the memory-bound regime: the weight matrix is PRE-QUANTIZED ONCE
3// (offline), and the HOT kernel streams ONLY the packed codes + per-block scales from DRAM -- the
4// exact deployment structure of ggml (weights quantized at load, dequant+dot in the inner loop).
5// This is the layer that, moving 3.5x less data at bounded quality, beats llama.cpp where it is
6// actually bottlenecked. The ENGINEER verifies the WHOLE layer output 1:1 vs the exact layer, and
7// the max per-row error is measured against the Council quality floor. license_tier: ORIGINAL
8
9import "nx_qmatvec.nx"     // QMV_B, QMV_Q4MAX/Q8MAX, qmv_block_scale, qmv_code, qmv_abs
10import "nx_vecmath.nx"
11
12func ql_blocks_per_row(c: i64) -> i64 { return (c + QMV_B - 1) / QMV_B }
13
14// OFFLINE (once, at load): quantize the R x C weight matrix into packed codes[R*C] + per-block
15// scales[R*nbpr]. Not on the hot path -- amortized over every token.
16func qlayer_quantize(w: *i64, r: i64, c: i64, qmax: i64, codes: *i64, scales: *i64) -> i64 {
17    let nbpr: i64 = ql_blocks_per_row(c)
18    var row: i64 = 0
19    while row < r {
20        let base: i64 = row * c
21        var bk: i64 = 0
22        while bk < nbpr {
23            let off: i64 = bk * QMV_B
24            var len: i64 = QMV_B; if off + len > c { len = c - off }
25            let s: i64 = qmv_block_scale(w, base + off, len)
26            scales[row * nbpr + bk] = s
27            var i: i64 = 0
28            while i < len { codes[base + off + i] = qmv_code(w[base + off + i], s, qmax); i = i + 1 }
29            bk = bk + 1
30        }
31        row = row + 1
32    }
33    return 0
34}
35
36// THE HOT KERNEL: y = dequant(codes,scales) . x, per row. Reads ONLY codes + scales + x (the data
37// streamed from DRAM in deployment). Scale factored out of the inner loop (ggml's trick).
38func qlayer_matvec(codes: *i64, scales: *i64, x: *i64, r: i64, c: i64, qmax: i64, y: *i64) -> i64 {
39    let nbpr: i64 = ql_blocks_per_row(c)
40    var row: i64 = 0
41    while row < r {
42        let base: i64 = row * c
43        var acc: i64 = 0; var bk: i64 = 0
44        while bk < nbpr {
45            let off: i64 = bk * QMV_B
46            var len: i64 = QMV_B; if off + len > c { len = c - off }
47            let s: i64 = scales[row * nbpr + bk]
48            var bdot: i64 = 0; var i: i64 = 0
49            while i < len { bdot = bdot + codes[base + off + i] * x[off + i]; i = i + 1 }
50            acc = acc + s * bdot
51            bk = bk + 1
52        }
53        y[row] = acc / qmax
54        row = row + 1
55    }
56    return 0
57}
58
59// the EXACT reference layer (the Engineer's oracle).
60func qlayer_matvec_exact(w: *i64, x: *i64, r: i64, c: i64, y: *i64) -> i64 {
61    var row: i64 = 0
62    while row < r {
63        let base: i64 = row * c
64        var acc: i64 = 0; var i: i64 = 0
65        while i < c { acc = acc + w[base + i] * x[i]; i = i + 1 }
66        y[row] = acc
67        row = row + 1
68    }
69    return 0
70}
71
72// integer floor sqrt (Newton). For the L2 norm without floats.
73func qlayer_isqrt(n: i64) -> i64 { return vm_isqrt(n) }
74
75// L2 RELATIVE error in PERMIL = sqrt( sum (yq-ye)^2 / sum ye^2 ) * 1000. This is the STANDARD
76// quantization-quality metric over the whole output vector -- robust to individual rows whose
77// exact value is near zero (where per-row RELATIVE error is meaningless). Honest measure of how
78// much the quantized layer's output deviates from the exact layer. (Per-row relative error is the
79// wrong tool here: a true-near-zero row makes it explode without any real quality loss.)
80func qlayer_l2_relerr_permil(yq: *i64, ye: *i64, r: i64) -> i64 {
81    var sd2: i64 = 0; var se2: i64 = 0; var row: i64 = 0
82    while row < r {
83        let d: i64 = yq[row] - ye[row]
84        sd2 = sd2 + d * d
85        se2 = se2 + ye[row] * ye[row]
86        row = row + 1
87    }
88    if se2 <= 0 { return 0 }
89    // permil = sqrt(sd2/se2)*1000 = sqrt(sd2*1e6/se2). guard the *1e6 against overflow by scaling
90    // both sums down equally while sd2 is large (preserves the ratio).
91    var a: i64 = sd2; var b: i64 = se2
92    let LIM: i64 = 9223372036854775
93    var g: i64 = 1
94    while a > LIM { a = a >> 1; b = b >> 1; g = g }   // keep a*1e6 inside i64
95    if b <= 0 { b = 1 }
96    return qlayer_isqrt((a * 1000000) / b)
97}
98
99// kept for diagnostics: MAX per-row relative error (UNSTABLE near zero -- use L2 for the verdict).
100func qlayer_max_err_permil(yq: *i64, ye: *i64, r: i64) -> i64 {
101    var mx: i64 = 0; var row: i64 = 0
102    while row < r {
103        let e: i64 = ye[row]
104        if e != 0 {
105            var d: i64 = yq[row] - e; if d < 0 { d = 0 - d }
106            var ea: i64 = e; if ea < 0 { ea = 0 - ea }
107            let rp: i64 = (d * 1000) / ea
108            if rp > mx { mx = rp }
109        }
110        row = row + 1
111    }
112    return mx
113}
114
115// total DRAM bytes the hot kernel streams for the weights, per format vs the int16 baseline.
116func qlayer_bytes_q(r: i64, c: i64, bits: i64) -> i64 { return (r * c * bits) / 8 + (r * ql_blocks_per_row(c)) * 2 }
117func qlayer_bytes_ref16(r: i64, c: i64) -> i64 { return r * c * 2 }
nx_qlayer.nx source

`nx_qlayer.nx` source
