`nx_qlayer.nx` source

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

`nx_qlayer.nx` source
