`nx_mcu_q4.nx` source

↩ module page · 104 lines · 4502 B
1// nx_mcu_q4.nx -- the int4-weight x int8-activation group dot product: the kernel MCU-class LLM inference
2// actually spends its time in. Pure integer, no float lib, no libc.
3//
4// WHY INTEGER AND WHY THIS KERNEL. On an ESP32-S3 the weights are group-wise symmetric int4 and the
5// activations are quantized to int8 per token, so every dot is int4 x int8 accumulated into int32. That
6// int32 group sum is EXACTLY what the S3 SIMD int8 dot instruction produces, which is why a scalar
7// reference and a vectorised device kernel are numerically identical rather than merely close. Being
8// integer, it is also bit-exact -- so its gate asserts EQUALITY, never a tolerance. A float kernel would
9// need an epsilon, and an epsilon is where a wrong answer hides.
10//
11// LAYOUT (mirrors the exported PLE1 format byte-for-byte):
12//   - a row is ceil(cols/2) bytes of packed nibbles; rows are byte-aligned, so odd cols leave a dead
13//     high nibble in the last byte. Ragged, no padding.
14//   - nibble = value + 8, giving the signed range -8..+7 from an unsigned 0..15 code.
15//   - EVEN column index -> LOW nibble; ODD column index -> HIGH nibble. Getting this backwards still
16//     produces plausible-looking numbers, so the gate pins it with a discriminating control.
17//   - activations are int8 two's complement in a byte, so they need sign extension, not a cast.
18//
19// Group scales are fp16 and are applied OUTSIDE this kernel (acc * scale), deliberately: this file stays
20// integer-exact and single-responsibility, and the float scaling lives with the code that owns floats.
21//
22// Semantics derived from the int8-activation path of esp32-ai firmware/common/llm.h (MIT, (c) 2026
23// Viacheslav Sierbov). Independent NishiLang implementation; behaviour is pinned by nx_mcu_q4_gate.
24// Shifts are written as / and % so this stays portable to the no-shift lowering path.
25// license_tier: DERIVED-MIT   No hw writes (Rule 26).
26import "nx_syscalls.nx"
27
28const Q4_BIAS: i64 = 8
29const Q4_CODES: i64 = 16
30const Q4_I8_WRAP: i64 = 256
31const Q4_I8_MAX: i64 = 127
32
33// bytes needed for one packed row of `cols` int4 values (ragged, byte-aligned)
34func q4_row_bytes(cols: i64) -> i64 {
35    if cols <= 0 { return 0 }
36    return (cols + 1) / 2
37}
38
39// sign-extend a byte holding an int8 two's complement activation: 0..127 stay, 128..255 become -128..-1
40func q4_i8(b: i64) -> i64 {
41    if b > Q4_I8_MAX { return b - Q4_I8_WRAP }
42    return b
43}
44
45// raw 0..15 nibble code at column j. EVEN j -> low nibble, ODD j -> high nibble.
46func q4_code_at(row: *u8, j: i64) -> i64 {
47    let byte: i64 = row[j / 2] as i64
48    if (j % 2) == 0 { return byte % Q4_CODES }
49    return byte / Q4_CODES
50}
51
52// signed weight value at column j, in -8..+7
53func q4_val_at(row: *u8, j: i64) -> i64 {
54    return q4_code_at(row, j) - Q4_BIAS
55}
56
57// THE HOT LOOP: int32-equivalent accumulator over columns [begin, end) of one row.
58// Returns sum( weight[j] * activation[j] ). Caller multiplies by the group scale.
59func q4_group_dot(row: *u8, xq: *u8, begin: i64, end: i64) -> i64 {
60    var g: i64 = 0
61    var j: i64 = begin
62    while j < end {
63        g = g + q4_val_at(row, j) * q4_i8(xq[j] as i64)
64        j = j + 1
65    }
66    return g
67}
68
69// full-row dot across every group boundary; scales are applied by the caller per group, so this is the
70// single-group / scale-free form used by the gate and by any caller with one group per row.
71func q4_row_dot(row: *u8, xq: *u8, cols: i64) -> i64 {
72    return q4_group_dot(row, xq, 0, cols)
73}
74
75// number of groups covering `cols` at group width `group` (last group is short, never padded)
76func q4_n_groups(cols: i64, group: i64) -> i64 {
77    if group <= 0 { return 0 }
78    if cols <= 0 { return 0 }
79    return (cols + group - 1) / group
80}
81
82// end column of group gi, clamped to cols -- the ragged-tail rule in one place
83func q4_group_end(cols: i64, group: i64, gi: i64) -> i64 {
84    var e: i64 = (gi + 1) * group
85    if e > cols { e = cols }
86    return e
87}
88
89// TEST/EXPORT HELPER: pack signed values (-8..+7) from vals[0..cols) into packed nibbles at dst.
90func q4_pack(dst: *u8, vals: *i64, cols: i64) -> i64 {
91    let rb: i64 = q4_row_bytes(cols)
92    var i: i64 = 0
93    while i < rb { dst[i] = 0 as u8; i = i + 1 }
94    var j: i64 = 0
95    while j < cols {
96        let code: i64 = vals[j] + Q4_BIAS
97        let bi: i64 = j / 2
98        var cur: i64 = dst[bi] as i64
99        if (j % 2) == 0 { cur = cur + code } else { cur = cur + code * Q4_CODES }
100        dst[bi] = cur as u8
101        j = j + 1
102    }
103    return 0
104}
nx_mcu_q4.nx source

`nx_mcu_q4.nx` source
