nx_quant_policy.nx source
↩ module page · 364 lines · 14999 B
1// nx_quant_policy.nx -- per-layer mixed-precision quantization policy.
2//
3// Ships VRAM-track Q-003 per docs/VRAM_OPTIMIZATION_REALISTIC_TRACKING.md.
4// Pure composition over the three quantizers already shipped:
5// nx_quant_block.nx -- q4_0 (4-bit symmetric, 10.67x compression)
6// nx_quant_block_q8.nx -- q8_0 (8-bit symmetric, 6.4x compression)
7// nx_quant_q4k.nx -- q4_K (4-bit k-quants w/ per-group scales)
8//
9// The model loader consults this policy to decide WHICH quantizer
10// to apply per-layer. Not every layer is equally sensitive to
11// quantization loss; cheap-to-compress layers (FFN up/down/gate)
12// go aggressive (q4_K), expensive-to-compress layers (Q/K/V
13// projections, embeddings, final output) stay safer (q8_0 or
14// passthrough at f16).
15//
16// Sensitivity heuristic literature:
17// - Frantar 2022 _GPTQ_: per-channel scale, calibration-data-driven
18// - Lin et al. 2023 _AWQ_: activation-aware -- protect outlier
19// channels at higher precision
20// - Xiao et al. 2023 _SmoothQuant_: pre-shift activation outliers
21// into weights so weights can be quantized further without
22// activation-side accuracy loss
23// - ggml q-mix presets 2024: practitioner-tuned per-tensor-name
24// defaults for Llama/Mistral/Qwen
25//
26// The substrate ships the POLICY TABLE structure + a default
27// heuristic. Per-model calibration (the AWQ pass) is a separate
28// workstream that fills the table with measured sensitivities.
29//
30// VRAM impact composition:
31// * 70% of params at q4_K (4-bit + k-quant overhead): ~3.5x save
32// * 25% of params at q8_0: ~1.4x save
33// * 5% of params at f16 passthrough: ~0.5x save
34// * Aggregate: ~3.0x save
35//
36// Per the bounded-loop cardinal: every loop here uses the
37// LoopVerdict pattern from nx_loop.nx.
38//
39// genealogy_id: gptq_frantar_2022 + awq_lin_2023 + smoothquant_xiao_2023 +
40// ggml_quant_mix_presets_2024
41// lineage_id: substrate_quant_policy_v1
42
43// nx_safety_envelope:
44// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
45// sil_target: SIL1
46// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
47// verdict: NOT_YET_EVALUATED
48
49import "nx_syscalls.nx"
50import "nx_tier.nx"
51import "nx_loop.nx"
52
53// ===== Sealed-enum: QuantPolicyKind ===============================
54//
55// One tag per quantizer the policy can dispatch to. Match
56// docs/VRAM_OPTIMIZATION_REALISTIC_TRACKING.md scenario inputs.
57
58const NX_QP_F16: nx_int = 0 // passthrough (no quantization)
59const NX_QP_Q8_0: nx_int = 1 // 8-bit symmetric, ~6.4x compression
60const NX_QP_Q4_K: nx_int = 2 // 4-bit k-quant, ~8.5x compression (with super-block overhead)
61const NX_QP_Q4_0: nx_int = 3 // 4-bit symmetric legacy, ~10.7x compression
62const NX_QP_F32: nx_int = 4 // full precision (rare, debug)
63const NX_QP_N_KINDS: nx_int = 5
64
65func nx_qp_is_valid(k: nx_int) -> nx_int {
66 if k < 0 { return 0 }
67 if k >= NX_QP_N_KINDS { return 0 }
68 return 1
69}
70
71// ===== Sealed-enum: LayerRole ====================================
72//
73// Per-layer role tag used by the heuristic. Captures the standard
74// transformer-block decomposition. Caller tags each layer when
75// registering; policy lookup uses the role to default-route.
76
77const NX_QR_UNKNOWN: nx_int = 0
78const NX_QR_EMBED: nx_int = 1 // token / position embeddings
79const NX_QR_ATTN_Q: nx_int = 2 // Q projection
80const NX_QR_ATTN_K: nx_int = 3 // K projection
81const NX_QR_ATTN_V: nx_int = 4 // V projection
82const NX_QR_ATTN_O: nx_int = 5 // output projection (after attention)
83const NX_QR_FFN_UP: nx_int = 6 // up-projection
84const NX_QR_FFN_DOWN: nx_int = 7 // down-projection
85const NX_QR_FFN_GATE: nx_int = 8 // gate projection (Llama-style)
86const NX_QR_NORM: nx_int = 9 // LayerNorm / RMSNorm scale
87const NX_QR_OUTPUT: nx_int = 10 // final output projection / lm_head
88const NX_QR_BIAS: nx_int = 11 // bias vectors (always small, keep f16)
89const NX_QR_N_ROLES: nx_int = 12
90
91func nx_qr_is_valid(r: nx_int) -> nx_int {
92 if r < 0 { return 0 }
93 if r >= NX_QR_N_ROLES { return 0 }
94 return 1
95}
96
97// ===== Policy entry ===============================================
98//
99// Caller registers each layer by (role, depth_from_first, depth_from_last).
100// `depth_from_*` lets the heuristic differentiate "first attention" from
101// "middle attention" -- early and late layers are more sensitive.
102
103struct NxQuantPolicyEntry {
104 role: nx_int,
105 depth_from_first: nx_int,
106 depth_from_last: nx_int,
107 quant_kind: nx_int // resolved by select_kind; NX_QP_*
108}
109
110const NX_QP_ENTRY_BYTES: nx_int = 32 // 4 fields * 8
111
112// ===== Default sensitivity heuristic ==============================
113//
114// Returns the recommended quant kind for a (role, depth_from_first,
115// depth_from_last) triple. Rules below crib from ggml practice +
116// AWQ outlier-channel protection.
117//
118// embeddings, lm_head, biases -> f16 (small + sensitive)
119// norm scales -> f16 (tiny tensors)
120// first 2 + last 2 attention K/V -> q8_0 (boundary layers sensitive)
121// other attention Q/K/V -> q8_0
122// attention O projection -> q4_K
123// FFN up/down/gate -> q4_K (most compressible)
124// unknown -> q8_0 (safe default)
125
126func nx_qp_select_kind(role: nx_int, depth_from_first: nx_int, depth_from_last: nx_int) -> nx_int {
127 if role == NX_QR_EMBED { return NX_QP_F16 }
128 if role == NX_QR_OUTPUT { return NX_QP_F16 }
129 if role == NX_QR_BIAS { return NX_QP_F16 }
130 if role == NX_QR_NORM { return NX_QP_F16 }
131
132 // Boundary layers (first 2 / last 2) keep K and V projections at q8_0.
133 if role == NX_QR_ATTN_K {
134 if depth_from_first < 2 { return NX_QP_Q8_0 }
135 if depth_from_last < 2 { return NX_QP_Q8_0 }
136 return NX_QP_Q8_0
137 }
138 if role == NX_QR_ATTN_V {
139 if depth_from_first < 2 { return NX_QP_Q8_0 }
140 if depth_from_last < 2 { return NX_QP_Q8_0 }
141 return NX_QP_Q8_0
142 }
143
144 if role == NX_QR_ATTN_Q { return NX_QP_Q8_0 }
145 if role == NX_QR_ATTN_O { return NX_QP_Q4_K }
146
147 if role == NX_QR_FFN_UP { return NX_QP_Q4_K }
148 if role == NX_QR_FFN_DOWN { return NX_QP_Q4_K }
149 if role == NX_QR_FFN_GATE { return NX_QP_Q4_K }
150
151 return NX_QP_Q8_0
152}
153
154// ===== Policy registry ============================================
155//
156// Caller allocates and fills. Bytes-per-param accounting works off
157// the resolved kinds; the lookup walks the entry array linearly --
158// fine for N_layers <= a few hundred which is the realistic scale.
159
160struct NxQuantPolicy {
161 n_entries: nx_int,
162 cap: nx_int,
163 entries: *NxQuantPolicyEntry
164}
165
166const NX_QP_REG_BYTES: nx_int = 24
167
168func nx_qp_alloc(cap: nx_int) -> *NxQuantPolicy {
169 let p: *NxQuantPolicy = sys_mmap(NX_QP_REG_BYTES) as *NxQuantPolicy
170 p.n_entries = 0
171 p.cap = cap
172 p.entries = sys_mmap(cap * NX_QP_ENTRY_BYTES) as *NxQuantPolicyEntry
173 return p
174}
175
176func nx_qp_register(p: *NxQuantPolicy,
177 role: nx_int,
178 depth_from_first: nx_int,
179 depth_from_last: nx_int) -> nx_int {
180 if p.n_entries >= p.cap { return 0 - 1 }
181 let idx: nx_int = p.n_entries
182 let e: *NxQuantPolicyEntry = (p.entries as i64 + idx * NX_QP_ENTRY_BYTES) as *NxQuantPolicyEntry
183 e.role = role
184 e.depth_from_first = depth_from_first
185 e.depth_from_last = depth_from_last
186 e.quant_kind = nx_qp_select_kind(role, depth_from_first, depth_from_last)
187 p.n_entries = p.n_entries + 1
188 return idx
189}
190
191// ===== Bytes-per-param table ======================================
192//
193// Used to compute aggregate VRAM under the policy. Numbers are
194// Q14 fixed point bits-per-param (16384 = 1 byte/param == 8 bpp):
195// f16 = 2 bytes/param -> 32768
196// q8_0 = 1 byte/param + scale overhead per 32 params
197// = 1 + 8/32 = 1.25 bytes/param -> 20480
198// q4_K = ~0.55 bytes/param including super-block scales
199// (4 bits + ~3.6% overhead) -> 9011
200// q4_0 = 0.5 bytes/param + scale overhead per 32 params
201// = 0.5 + 8/32 = 0.75 bytes/param -> 12288
202// f32 = 4 bytes/param -> 65536
203
204func nx_qp_bytes_per_param_q14(kind: nx_int) -> nx_int {
205 if kind == NX_QP_F16 { return 32768 }
206 if kind == NX_QP_Q8_0 { return 20480 }
207 if kind == NX_QP_Q4_K { return 9011 }
208 if kind == NX_QP_Q4_0 { return 12288 }
209 if kind == NX_QP_F32 { return 65536 }
210 return 32768 // safe default
211}
212
213// ===== Aggregate VRAM accounting ==================================
214//
215// Given the policy + per-entry parameter counts, return the total
216// quantized byte budget in Q0 (raw bytes). Honest only after caller
217// passes real counts -- substrate has no idea what the model weights
218// look like until told.
219
220func nx_qp_total_bytes(p: *NxQuantPolicy, params_per_entry: *i64) -> nx_int {
221 let n: nx_int = p.n_entries
222 var iter: nx_int = 0
223 var verdict: nx_int = NX_LOOP_RUNNING
224 let BUDGET: nx_int = n
225 var total_q14: i64 = 0
226 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
227 let e: *NxQuantPolicyEntry = (p.entries as i64 + iter * NX_QP_ENTRY_BYTES) as *NxQuantPolicyEntry
228 let bpp: nx_int = nx_qp_bytes_per_param_q14(e.quant_kind)
229 total_q14 = total_q14 + params_per_entry[iter] * bpp
230 iter = iter + 1
231 }
232 // Convert Q14 byte-count to raw byte-count.
233 return total_q14 / 16384
234}
235
236func nx_qp_total_bytes_dense_f16(params_per_entry: *i64, n: nx_int) -> nx_int {
237 var iter: nx_int = 0
238 var verdict: nx_int = NX_LOOP_RUNNING
239 let BUDGET: nx_int = n
240 var total: i64 = 0
241 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
242 total = total + params_per_entry[iter] * 2
243 iter = iter + 1
244 }
245 return total
246}
247
248// ===== Compression ratio (Q10) ====================================
249
250func nx_qp_compression_ratio_q10(p: *NxQuantPolicy, params_per_entry: *i64) -> nx_int {
251 let dense: nx_int = nx_qp_total_bytes_dense_f16(params_per_entry, p.n_entries)
252 let mixed: nx_int = nx_qp_total_bytes(p, params_per_entry)
253 if mixed <= 0 { return 0 }
254 return (dense * 1024) / mixed
255}
256
257// ===== Self-test ==================================================
258//
259// Three closed-form invariants:
260//
261// (a) Role-routing -- every role maps to the documented kind.
262// (b) Aggregate compression ratio on a synthetic 16-layer
263// transformer (1 embed + 14 blocks * 9 sublayers + 1 output +
264// 1 lm_head) is within ±0.05x of the documented 2.5x-3.5x band
265// per the doc's Q-003 scenario.
266// (c) Bytes-per-param table is monotone: f32 > f16 > q8_0 > q4_0
267// > q4_K (q4_K beats q4_0 due to better packing despite shared
268// 4-bit width).
269
270func main() -> i64 {
271 // --- (a) Role routing ---
272 if nx_qp_select_kind(NX_QR_EMBED, 0, 31) != NX_QP_F16 { return 10 }
273 if nx_qp_select_kind(NX_QR_OUTPUT, 31, 0) != NX_QP_F16 { return 11 }
274 if nx_qp_select_kind(NX_QR_BIAS, 5, 25) != NX_QP_F16 { return 12 }
275 if nx_qp_select_kind(NX_QR_NORM, 5, 25) != NX_QP_F16 { return 13 }
276 if nx_qp_select_kind(NX_QR_ATTN_Q, 5, 25) != NX_QP_Q8_0 { return 14 }
277 if nx_qp_select_kind(NX_QR_ATTN_K, 5, 25) != NX_QP_Q8_0 { return 15 }
278 if nx_qp_select_kind(NX_QR_ATTN_V, 5, 25) != NX_QP_Q8_0 { return 16 }
279 if nx_qp_select_kind(NX_QR_ATTN_O, 5, 25) != NX_QP_Q4_K { return 17 }
280 if nx_qp_select_kind(NX_QR_FFN_UP, 5, 25) != NX_QP_Q4_K { return 18 }
281 if nx_qp_select_kind(NX_QR_FFN_DOWN, 5, 25) != NX_QP_Q4_K { return 19 }
282 if nx_qp_select_kind(NX_QR_FFN_GATE, 5, 25) != NX_QP_Q4_K { return 20 }
283 if nx_qp_select_kind(NX_QR_UNKNOWN, 5, 25) != NX_QP_Q8_0 { return 21 }
284
285 // --- (b) Monotone bytes-per-param ---
286 let bf32: nx_int = nx_qp_bytes_per_param_q14(NX_QP_F32)
287 let bf16: nx_int = nx_qp_bytes_per_param_q14(NX_QP_F16)
288 let bq8: nx_int = nx_qp_bytes_per_param_q14(NX_QP_Q8_0)
289 let bq4_0: nx_int = nx_qp_bytes_per_param_q14(NX_QP_Q4_0)
290 let bq4_K: nx_int = nx_qp_bytes_per_param_q14(NX_QP_Q4_K)
291 if bf32 <= bf16 { return 30 }
292 if bf16 <= bq8 { return 31 }
293 if bq8 <= bq4_0 { return 32 }
294 if bq4_0 <= bq4_K { return 33 }
295
296 // --- (c) Synthetic 16-layer compression budget ---
297 // Layout: 1 embed + 14*(Q,K,V,O,FFN_UP,FFN_DOWN,FFN_GATE,NORM1,NORM2)
298 // + 1 output + 1 lm_head.
299 // n_entries = 1 + 14*9 + 1 + 1 = 129.
300 let cap: nx_int = 256
301 let pol: *NxQuantPolicy = nx_qp_alloc(cap)
302
303 // Embed: large (1M params), depth 0.
304 nx_qp_register(pol, NX_QR_EMBED, 0, 15)
305
306 // 14 transformer blocks.
307 var blk: nx_int = 0
308 var verdict: nx_int = NX_LOOP_RUNNING
309 let N_BLK: nx_int = 14
310 while verdict == NX_LOOP_RUNNING && blk < N_BLK {
311 let df: nx_int = blk
312 let dl: nx_int = 15 - blk
313 nx_qp_register(pol, NX_QR_ATTN_Q, df, dl)
314 nx_qp_register(pol, NX_QR_ATTN_K, df, dl)
315 nx_qp_register(pol, NX_QR_ATTN_V, df, dl)
316 nx_qp_register(pol, NX_QR_ATTN_O, df, dl)
317 nx_qp_register(pol, NX_QR_FFN_UP, df, dl)
318 nx_qp_register(pol, NX_QR_FFN_DOWN, df, dl)
319 nx_qp_register(pol, NX_QR_FFN_GATE, df, dl)
320 nx_qp_register(pol, NX_QR_NORM, df, dl)
321 nx_qp_register(pol, NX_QR_NORM, df, dl)
322 blk = blk + 1
323 }
324 nx_qp_register(pol, NX_QR_OUTPUT, 15, 0)
325 nx_qp_register(pol, NX_QR_BIAS, 15, 0)
326 if pol.n_entries != 129 { return 40 }
327
328 // Param counts: 1M embed + 1M output + tiny bias + 0.5M per
329 // attn sublayer + 1.0M per FFN sublayer + tiny norms.
330 let pp: *i64 = sys_mmap(pol.n_entries * 8) as *i64
331 pp[0] = 1000000 // embed
332 var idx: nx_int = 1
333 var bi: nx_int = 0
334 while bi < N_BLK {
335 pp[idx + 0] = 500000 // Q
336 pp[idx + 1] = 500000 // K
337 pp[idx + 2] = 500000 // V
338 pp[idx + 3] = 500000 // O
339 pp[idx + 4] = 1000000 // FFN_UP
340 pp[idx + 5] = 1000000 // FFN_DOWN
341 pp[idx + 6] = 1000000 // FFN_GATE
342 pp[idx + 7] = 1024 // NORM
343 pp[idx + 8] = 1024 // NORM
344 idx = idx + 9
345 bi = bi + 1
346 }
347 pp[idx] = 1000000 // output
348 pp[idx + 1] = 8192 // bias
349
350 // Total params: 1M (embed) + 14 * (4*0.5M + 3*1M + 2*1K) + 1M + 8K
351 // = 1M + 14*5_002_048 + 1_008_192
352 // = 1M + 70_028_672 + 1_008_192
353 // = 72_036_864 params total.
354 // Dense f16: 144_073_728 bytes (~137 MB).
355 // Under the policy: most params are FFN/attn at q4_K/q8_0;
356 // embed + output stay f16. Expected ratio ~2.5..3.5x per
357 // the Q-003 row of the tracking doc.
358 let r: nx_int = nx_qp_compression_ratio_q10(pol, pp)
359 // 2.5x -> 2560 Q10; 3.5x -> 3584 Q10. Allow 2400..3700.
360 if r < 2400 { return 50 }
361 if r > 3700 { return 51 }
362
363 return 0
364}