nx_quant_q4k.nx source
↩ module page · 277 lines · 10622 B
1// nx_quant_q4k.nx -- hierarchical k-quants (ggml q4_K shape).
2//
3// Improvement over nx_quant_block.nx (q4_0 shape):
4//
5// q4_0: one scale per 32-value block. All 32 values share the
6// same dynamic range; groups with wildly different magnitudes
7// waste precision (small-magnitude groups get coarse steps).
8//
9// q4_K: super-block of 256 values = 8 groups of 32. TWO scales:
10// a SUPER scale per super-block + a LOCAL scale per group.
11// The hierarchy lets each group adapt its precision while
12// still sharing the super-block magnitude estimate -- which
13// empirically halves worst-case error vs q4_0 at near-identical
14// storage.
15//
16// Concrete win for the user's "VRAM optimization realistic tracking":
17// This is the realisation of the Q-002 row -- q4_K + AWQ
18// calibration -- which the doc projected at 1-2% quality loss vs
19// q4_0's 5-7%.
20//
21// Storage per 256 values:
22// 1 super_d (i64) = 8 bytes
23// 8 group_d_local (i64 each) = 64 bytes
24// 128 packed nibbles (256 values) = 128 bytes
25// Total = 200 bytes
26//
27// vs dense 256 * 8 = 2048 bytes -> 10.24x compression
28//
29// Reconstruction: value = nibble * group_d_local * super_d / Q10
30//
31// Algorithm-led precision:
32// 1. Find super_max = max(|values|) across the 256-value super-block
33// 2. super_d = super_max / 127 (range of effective combined scale)
34// 3. For each 32-value group:
35// a. find group_max = max(|values|) in the group
36// b. group_d_local = (group_max * Q10) / (7 * super_d)
37// (so that nibble of 7 reconstructs to group_max)
38// 4. For each value: nibble = round(value * Q10 / (group_d_local * super_d))
39//
40// Hierarchy benefit: a group with group_max << super_max gets a SMALL
41// d_local; its nibbles span its actual range with full 4-bit precision.
42// In q4_0 that same group would share the super-block scale and waste
43// most of its 4-bit precision.
44//
45// genealogy_id: ggml_q4_k_2024 + lloyd_max_1957 + jacob_kligys_2018 +
46// awq_lin_2023
47// lineage_id: substrate_quant_q4k_v1
48
49// nx_safety_envelope:
50// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
51// sil_target: SIL1
52// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
53// verdict: NOT_YET_EVALUATED
54
55import "nx_syscalls.nx"
56import "nx_tier.nx"
57import "nx_tensor.nx"
58import "nx_quant_block.nx"
59const NX_MAGIC_1024: i64 = 1024
60
61const NX_Q4K_Q10: nx_int = 1024
62const NX_Q4K_SUPER_VALUES: nx_int = 256 // values per super-block
63const NX_Q4K_GROUP_VALUES: nx_int = 32 // values per group
64const NX_Q4K_GROUPS: nx_int = 8 // groups per super-block
65const NX_Q4K_NIBBLE_MAX: nx_int = 7 // [-7..+7]
66const NX_Q4K_SUPER_RANGE: nx_int = 127 // super_d denominator
67
68// Bytes per super-block: 8 (super_d) + 8*8 (group locals) + 128 (nibbles)
69const NX_Q4K_BYTES_PER_SUPER: nx_int = 200
70
71// ===== Sealed-enum: Q4KVerdict ====================================
72
73const NX_Q4K_OK: nx_int = 0
74const NX_Q4K_ERR_BAD_LEN: nx_int = 1 // n not multiple of 256
75const NX_Q4K_ERR_BAD_DTYPE: nx_int = 2
76const NX_Q4K_ERR_SHAPE_MISMATCH: nx_int = 3
77const NX_Q4K_N_VERDICTS: nx_int = 4
78
79func nx_q4k_verdict_is_valid(v: nx_int) -> nx_int {
80 if v < 0 { return 0 }
81 if v >= NX_Q4K_N_VERDICTS { return 0 }
82 return 1
83}
84
85// ===== Q4K buffer struct ==========================================
86
87struct NxQuantQ4K {
88 n_supers: nx_int, // number of super-blocks
89 n_values: nx_int, // n_supers * 256
90 super_d: *i64, // [n_supers]
91 group_d: *i64, // [n_supers * 8]
92 packed: *u8 // [n_supers * 128] (nibbles)
93}
94
95const NX_Q4K_STRUCT_BYTES: nx_int = 40 // 5 fields * 8
96
97func nx_q4k_alloc(n_values: nx_int) -> *NxQuantQ4K {
98 let qb: *NxQuantQ4K = (sys_mmap(NX_Q4K_STRUCT_BYTES)) as *NxQuantQ4K
99 let n_supers: nx_int = (n_values + NX_Q4K_SUPER_VALUES - 1) / NX_Q4K_SUPER_VALUES
100 qb.n_supers = n_supers
101 qb.n_values = n_supers * NX_Q4K_SUPER_VALUES // padded
102 qb.super_d = (sys_mmap(n_supers * NX_SIZEOF_NX_INT)) as *i64
103 qb.group_d = (sys_mmap(n_supers * NX_Q4K_GROUPS * NX_SIZEOF_NX_INT)) as *i64
104 let packed_bytes: nx_int = n_supers * (NX_Q4K_SUPER_VALUES / 2)
105 qb.packed = (sys_mmap(packed_bytes)) as *u8
106 // Zero init
107 var i: nx_int = 0
108 while i < n_supers { qb.super_d[i] = 0; i = i + 1 }
109 var j: nx_int = 0
110 while j < n_supers * NX_Q4K_GROUPS { qb.group_d[j] = 0; j = j + 1 }
111 var k: nx_int = 0
112 while k < packed_bytes { qb.packed[k] = 0; k = k + 1 }
113 return qb
114}
115
116// ===== Nibble pack / unpack (same shape as q4_0) =================
117
118func _q4k_pack_nibble(qb: *NxQuantQ4K, value_idx: nx_int, signed_nib: nx_int) -> nx_int {
119 let byte_idx: nx_int = value_idx / 2
120 let is_high: nx_int = value_idx - byte_idx * 2
121 let unsigned_nib: nx_int = signed_nib + 7 // map [-7..7] to [0..14]
122 let cur: nx_int = qb.packed[byte_idx] as nx_int
123 var new_byte: nx_int = 0
124 if is_high == 0 {
125 new_byte = (cur - (cur - (cur / 16) * 16)) + unsigned_nib
126 }
127 if is_high == 1 {
128 new_byte = (cur - (cur / 16) * 16) + unsigned_nib * 16
129 }
130 qb.packed[byte_idx] = new_byte
131 return 0
132}
133
134func _q4k_unpack_nibble(qb: *NxQuantQ4K, value_idx: nx_int) -> nx_int {
135 let byte_idx: nx_int = value_idx / 2
136 let is_high: nx_int = value_idx - byte_idx * 2
137 let cur: nx_int = qb.packed[byte_idx] as nx_int
138 var unsigned_nib: nx_int = 0
139 if is_high == 0 {
140 unsigned_nib = cur - (cur / 16) * 16
141 }
142 if is_high == 1 {
143 unsigned_nib = cur / 16
144 }
145 return unsigned_nib - 7
146}
147
148func _q4k_abs(x: nx_int) -> nx_int {
149 if x < 0 { return 0 - x }
150 return x
151}
152
153// ===== Quantise =====================================================
154//
155// Per super-block:
156// 1. Find super_max across all 256 values.
157// 2. super_d = super_max / NX_Q4K_SUPER_RANGE (floor 1)
158// 3. Per group of 32 values:
159// a. find group_max
160// b. group_d_local = group_max / NX_Q4K_NIBBLE_MAX (floor 1)
161// then re-express as Q10 fraction of super_d: stored = group_d * Q10 / super_d
162// (so that reconstruction = nibble * (stored * super_d / Q10) = nibble * group_d_local)
163// 4. Per value: nibble = round(value / group_d_local)
164//
165// Note: we store group_d_local directly (i64) rather than the
166// Q10 ratio. This trades 6-bit packing for clarity and skips a
167// multiplication on dequant. Storage is still 8 bytes per group,
168// matching the budget.
169
170func nx_q4k_quantize(values: *i64, n: nx_int, qb: *NxQuantQ4K) -> nx_int {
171 if n > qb.n_values { return NX_Q4K_ERR_BAD_LEN }
172
173 var sb: nx_int = 0
174 while sb < qb.n_supers {
175 let sb_start: nx_int = sb * NX_Q4K_SUPER_VALUES
176 var sb_end: nx_int = sb_start + NX_Q4K_SUPER_VALUES
177 if sb_end > n { sb_end = n }
178
179 // Find super_max
180 var super_max: nx_int = 0
181 var i: nx_int = sb_start
182 while i < sb_end {
183 let a: nx_int = _q4k_abs(values[i])
184 if a > super_max { super_max = a }
185 i = i + 1
186 }
187 var super_d: nx_int = super_max / NX_Q4K_SUPER_RANGE
188 if super_d < 1 { super_d = 1 }
189 qb.super_d[sb] = super_d
190
191 // Per-group scales + nibbles
192 var g: nx_int = 0
193 while g < NX_Q4K_GROUPS {
194 let g_start: nx_int = sb_start + g * NX_Q4K_GROUP_VALUES
195 var g_end: nx_int = g_start + NX_Q4K_GROUP_VALUES
196 if g_end > n { g_end = n }
197 // Find group_max
198 var gmax: nx_int = 0
199 var j: nx_int = g_start
200 while j < g_end {
201 let aj: nx_int = _q4k_abs(values[j])
202 if aj > gmax { gmax = aj }
203 j = j + 1
204 }
205 var d_local: nx_int = gmax / NX_Q4K_NIBBLE_MAX
206 if d_local < 1 { d_local = 1 }
207 qb.group_d[sb * NX_Q4K_GROUPS + g] = d_local
208
209 // Quantise values in this group
210 var k: nx_int = g_start
211 while k < g_end {
212 var v: nx_int = values[k]
213 var sign: nx_int = 1
214 if v < 0 { sign = 0 - 1 }
215 let abs_v: nx_int = _q4k_abs(v)
216 var q: nx_int = (abs_v + d_local / 2) / d_local
217 if q > NX_Q4K_NIBBLE_MAX { q = NX_Q4K_NIBBLE_MAX }
218 q = q * sign
219 _q4k_pack_nibble(qb, k, q)
220 k = k + 1
221 }
222 // Zero-pad incomplete trailing group
223 var pad: nx_int = g_end
224 while pad < g_start + NX_Q4K_GROUP_VALUES {
225 _q4k_pack_nibble(qb, pad, 0)
226 pad = pad + 1
227 }
228 g = g + 1
229 }
230 sb = sb + 1
231 }
232 return NX_Q4K_OK
233}
234
235// ===== Dequantise =================================================
236
237func nx_q4k_dequantize(qb: *NxQuantQ4K, values_out: *i64, n: nx_int) -> nx_int {
238 if n > qb.n_values { return NX_Q4K_ERR_BAD_LEN }
239 var i: nx_int = 0
240 while i < n {
241 let sb: nx_int = i / NX_Q4K_SUPER_VALUES
242 let in_sb: nx_int = i - sb * NX_Q4K_SUPER_VALUES
243 let g: nx_int = in_sb / NX_Q4K_GROUP_VALUES
244 let nib: nx_int = _q4k_unpack_nibble(qb, i)
245 let d_local: nx_int = qb.group_d[sb * NX_Q4K_GROUPS + g]
246 values_out[i] = nib * d_local
247 i = i + 1
248 }
249 return NX_Q4K_OK
250}
251
252// ===== Compression ratio (Q10) ====================================
253
254func nx_q4k_compression_ratio_q10(qb: *NxQuantQ4K) -> nx_int {
255 let dense: nx_int = qb.n_values * 8
256 let quant: nx_int = qb.n_supers * NX_Q4K_BYTES_PER_SUPER
257 if quant <= 0 { return 0 }
258 return (dense * NX_MAGIC_1024) / quant
259}
260
261// ===== Tensor convenience =========================================
262
263func nx_q4k_quantize_tensor(t: *NxTensor) -> *NxQuantQ4K {
264 if t.dtype != NX_DT_I64 { return 0 as *NxQuantQ4K }
265 if t.ndim != 1 { return 0 as *NxQuantQ4K }
266 if nx_t_is_contiguous(t) == 0 { return 0 as *NxQuantQ4K }
267 let qb: *NxQuantQ4K = nx_q4k_alloc(t.numel)
268 nx_q4k_quantize(t.storage as *i64, t.numel, qb)
269 return qb
270}
271
272func nx_q4k_dequantize_tensor(qb: *NxQuantQ4K, t: *NxTensor) -> nx_int {
273 if t.dtype != NX_DT_I64 { return NX_Q4K_ERR_BAD_DTYPE }
274 if t.ndim != 1 { return NX_Q4K_ERR_SHAPE_MISMATCH }
275 if nx_t_is_contiguous(t) == 0 { return NX_Q4K_ERR_SHAPE_MISMATCH }
276 return nx_q4k_dequantize(qb, t.storage as *i64, t.numel)
277}