nx_gen_lora_repr.nx source
↩ module page · 179 lines · 8762 B
1// nx_gen_lora_repr.nx -- DECIDE the adapter representation by measuring it, not by arguing it.
2//
3// The hot matmul consumes Q8_0 blocks + per-block f16 scales. nx_lora_fold produces f32. So an
4// adapted run has exactly two shapes, and they trade accuracy against per-step cost:
5//
6// A) FOLD-AND-REQUANTIZE — dequantize Q8_0, fold, requantize to Q8_0, run the normal kernel.
7// One-time cost, zero per-step cost, hot loop untouched. This is what a "merged checkpoint"
8// is. But it quantizes a weight that was ALREADY quantized once.
9//
10// B) RANK-FACTORED DELTA — keep the original Q8_0 blocks untouched and add
11// scale * (h @ down^T) @ up^T
12// in f32 at matmul time. The base carries NO new error. Costs 2*rank*(in+out) per token per
13// step forever -- what sd.cpp pays.
14//
15// ★ BOTH ARE GRADED AGAINST THE EXACT f64 REFERENCE, NOT AGAINST EACH OTHER. Comparing A to B
16// answers "do they differ", which is the question nobody asked; comparing each to the truth answers
17// "which is right". This is the same two-way-comparison trap that made the oracle's quantization
18// error look like ours for most of this lane.
19//
20// Usage: nx_gen_lora_repr <gguf> <lora.safetensors> [multiplier_milli]
21// license_tier: ORIGINAL
22
23import "nx_syscalls.nx"
24import "nx_le.nx"
25import "nx_f32.nx"
26import "nx_f32_div.nx"
27import "nx_f32_cvt.nx"
28import "nx_f16.nx"
29import "nx_strconv.nx"
30import "nx_genfix.nx"
31import "nx_genver.nx"
32import "nx_genweights.nx"
33import "nx_genblock.nx"
34import "nx_q8_0_from_f32.nx"
35import "nx_genlora.nx"
36
37const LR_IN: i64 = 3840
38const LR_OUT: i64 = 11520
39const LR_TOK: i64 = 768
40
41func lr_puts(s: *u8) -> i64 {
42 var n: i64 = 0
43 while s[n] != (0 as u8) { n = n + 1 }
44 return sys_write(1, s, n)
45}
46
47// Quantize f32 [out_dim, in_dim] -> Q8_0 blocks, ROW BY ROW, using the estate's own quantizer.
48//
49// ⚠I had written my own absmax/round loop here before checking. `nx_q8_0_from_f32` already exists
50// -- the standard ggml Q8_0 quantizer, in production use for re-quantizing Q6_K at load. Writing a
51// second one would have put two rounding rules in the tree, and a rounding difference reads as
52// "quantization is just lossy" rather than as a bug.
53// ★ THE PRIMITIVE YOU ARE ABOUT TO WRITE IS THE ONE TO GREP FOR FIRST.
54// Its input is an i64 array of f32 BITS (not packed bytes), so feed it a row at a time -- 30KB of
55// scratch instead of a 354MB whole-tensor copy.
56func lr_quant_q8(src: *u8, dst: *u8, in_dim: i64, out_dim: i64, scales: *i64) -> i64 {
57 let nblk: i64 = in_dim / 32
58 let row: *i64 = sys_mmap(in_dim * 8 + 64) as *i64
59 var o: i64 = 0
60 while o < out_dim {
61 var i: i64 = 0
62 while i < in_dim { row[i] = nx_le_read_u32(src, (o * in_dim + i) * 4); i = i + 1 }
63 nx_q8_0_from_f32(row, in_dim, ((dst as i64) + o * nblk * 34) as *u8)
64 // read the scales back out of the blocks we just wrote, so the kernel sees exactly the
65 // f16-rounded value it would see from a file -- not the f32 we happened to compute.
66 var b: i64 = 0
67 while b < nblk {
68 scales[o * nblk + b] = nx_f16_to_f32(nx_le_read_u16(dst, (o * nblk + b) * 34))
69 b = b + 1
70 }
71 o = o + 1
72 }
73 return 0
74}
75
76func main(argc: i64, argv: *i64) -> i64 {
77 if argc < 3 {
78 lr_puts("usage: nx_gen_lora_repr <gguf> <lora.safetensors> [multiplier_milli]\n" as *u8)
79 return 2
80 }
81 var mm: i64 = 1000
82 if argc > 3 { mm = nx_strconv_parse_i64(argv[3] as *u8, sys_mmap(64) as *i64) }
83 let mult: i64 = nx_f32_div(nx_i32_to_f32(mm), nx_i32_to_f32(1000))
84
85 let M: *u8 = "lora_probe" as *u8
86 let n_h: *u8 = "layers.0.h_attn" as *u8
87 let n_r: *u8 = "qkv_lora.exact" as *u8
88 let n_w: *u8 = "model.diffusion_model.layers.0.attention.qkv.weight" as *u8
89
90 let h: *u8 = nx_genfix_load(M, n_h, br_strlen(n_h), LR_TOK * LR_IN)
91 if (h as i64) == 0 { lr_puts("h_attn fixture missing\n" as *u8); return 10 }
92 let ref: *u8 = nx_genfix_load(M, n_r, br_strlen(n_r), LR_TOK * LR_OUT)
93 if (ref as i64) == 0 { lr_puts("exact reference missing\n" as *u8); return 11 }
94 let Wf: *u8 = nx_genfix_load(M, n_w, br_strlen(n_w), LR_IN * LR_OUT)
95 if (Wf as i64) == 0 { lr_puts("base weight fixture missing\n" as *u8); return 12 }
96
97 let gw: *i64 = nx_gw_open(argv[1] as *u8)
98 if (gw as i64) == 0 { lr_puts("gguf open failed\n" as *u8); return 13 }
99 if br_arch_bind(gw) != 0 { lr_puts("arch not derivable\n" as *u8); return 14 }
100 let lw: *i64 = nx_gw_open(argv[2] as *u8)
101 if (lw as i64) == 0 { lr_puts("adapter unreadable\n" as *u8); return 15 }
102
103 let nblk: i64 = LR_IN / 32
104 // ⚠⚠ br_matmul FORKS. Its children write the output band, so the output buffer MUST be a
105 // SHARED mapping. sys_mmap is MAP_PRIVATE (0x22); sys_mmap_shared is MAP_SHARED (0x21).
106 // With a private buffer every child's work is discarded at exit and the parent reads ZEROS --
107 // no error, no signal, just a plausible-looking all-zero tensor. It cost both arms of this
108 // experiment reporting an IDENTICAL 98.6% failure, which is what finally gave it away.
109 // ★★★★★★ TWO INDEPENDENT ARMS THAT FAIL IDENTICALLY ARE NOT TWO RESULTS -- THE THING YOU VARIED
110 // IS NOT IN THE PATH. A difference you cannot see is evidence about your harness, not the world.
111 let y: *u8 = sys_mmap_shared(LR_TOK * LR_OUT * 4 + 64)
112 let tol: *i64 = sys_mmap(64) as *i64
113 nx_genver_tols(tol)
114
115 // ================= A) FOLD, THEN REQUANTIZE =================
116 let Wa: *u8 = sys_mmap(LR_IN * LR_OUT * 4 + 64)
117 var i: i64 = 0
118 while i < LR_IN * LR_OUT { nx_le_write_u32(Wa, i * 4, nx_le_read_u32(Wf, i * 4)); i = i + 1 }
119 let t0: i64 = sys_now_us()
120 if nx_lora_fold(lw, n_w, Wa, LR_IN, LR_OUT, mult) != LORA_APPLIED {
121 lr_puts("adapter did not apply\n" as *u8); return 16
122 }
123 let qa: *u8 = sys_mmap(LR_OUT * nblk * 34 + 64)
124 let sa: *i64 = sys_mmap(LR_OUT * nblk * 8 + 64) as *i64
125 lr_quant_q8(Wa, qa, LR_IN, LR_OUT, sa)
126 nx_genver_emit("A_prep_us" as *u8, sys_now_us() - t0)
127
128 let t1: i64 = sys_now_us()
129 br_matmul(qa, sa, h, y, LR_TOK, LR_IN, LR_OUT, 8)
130 nx_genver_emit("A_step_us" as *u8, sys_now_us() - t1)
131 let ca: *i64 = sys_mmap(128) as *i64
132 nx_genver_init(ca, 3)
133 i = 0
134 while i < LR_TOK * LR_OUT {
135 nx_genver_tally(ca, tol, nx_le_read_u32(y, i * 4), nx_le_read_u32(ref, i * 4), i)
136 i = i + 1
137 }
138 lr_puts("== A) fold + requantize to Q8_0 (one-time prep, zero per-step cost)\n" as *u8)
139 nx_genver_report(ca)
140
141 // ================= B) ORIGINAL Q8_0 + RANK-FACTORED f32 DELTA =================
142 // The base keeps its ORIGINAL blocks, so it carries no new quantization error at all.
143 let t2: i64 = sys_now_us()
144 let sb: *i64 = sys_mmap(LR_OUT * nblk * 8 + 64) as *i64
145 let qb: *u8 = br_gw_q8(gw, n_w, LR_IN, LR_OUT, sb)
146 if (qb as i64) == 0 { lr_puts("base Q8_0 tensor not found\n" as *u8); return 17 }
147 nx_genver_emit("B_prep_us" as *u8, sys_now_us() - t2)
148
149 let t3: i64 = sys_now_us()
150 br_matmul(qb, sb, h, y, LR_TOK, LR_IN, LR_OUT, 8)
151 // delta = mult * (alpha/rank) * (h @ down^T) @ up^T, applied on top. Rebuilt here from the SAME
152 // fold path so the two arms differ ONLY in representation, never in the adapter maths.
153 let Wd: *u8 = sys_mmap(LR_IN * LR_OUT * 4 + 64)
154 i = 0
155 while i < LR_IN * LR_OUT { nx_le_write_u32(Wd, i * 4, 0); i = i + 1 }
156 if nx_lora_fold(lw, n_w, Wd, LR_IN, LR_OUT, mult) != LORA_APPLIED {
157 lr_puts("delta build failed\n" as *u8); return 18
158 }
159 let dy: *u8 = sys_mmap_shared(LR_TOK * LR_OUT * 4 + 64)
160 br_matmul_f32(Wd, h, dy, LR_TOK, LR_IN, LR_OUT, 8)
161 i = 0
162 while i < LR_TOK * LR_OUT {
163 nx_le_write_u32(y, i * 4, __f32_add(nx_le_read_u32(y, i * 4), nx_le_read_u32(dy, i * 4)))
164 i = i + 1
165 }
166 // ⚠B_step_us is an UPPER BOUND, not the rank-factored cost: this arm materialises the dense
167 // [out,in] delta and runs a full f32 matmul. Its ACCURACY is the honest number (same maths);
168 // its TIME is the worst case. Quoting it as "the runtime-adapter cost" would be a fabrication.
169 nx_genver_emit("B_step_us_DENSE_UPPER_BOUND" as *u8, sys_now_us() - t3)
170 let cb: *i64 = sys_mmap(128) as *i64
171 nx_genver_init(cb, 3)
172 i = 0
173 while i < LR_TOK * LR_OUT {
174 nx_genver_tally(cb, tol, nx_le_read_u32(y, i * 4), nx_le_read_u32(ref, i * 4), i)
175 i = i + 1
176 }
177 lr_puts("== B) original Q8_0 base + f32 adapter delta (base unquantized-again)\n" as *u8)
178 return nx_genver_report(cb)
179}