nx_q4k_dot_simd_lib.nx source
↩ module page · 146 lines · 7177 B
1// nx_q4k_dot_simd_lib.nx -- the fused Q4_K SIMD dot as a LIBRARY (extracted 2026-09-02 from nx_q4k_dot_simd.nx,
2// which keeps its bench main and now imports this; the extraction is proven by that organ rebuilding to the
3// banked sha). Consumers: nx_nofloat_q4k (LM4 resident-quant decode). Kernel unchanged: per sub-block
4// dot += d1*sq - m0*sc with sq = sum q*col via vpmaddwd (q i16 0..15, col i16), sc from a precomputed
5// activation-only table. Returns sum over the row of W_Q24 * col_i16 (weight in Q24 fixed point).
6// license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_le.nx"
10import "nx_tensor.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13
14// NIBBLE UNPACK BY 64-BIT SPREAD (2026-09-02, LM4b). MEASURED on the LM4 serve: the byte-at-a-time unpack (four
15// nx_le_read_u8 calls plus a pack per 4 bytes, redone for every output row on every token) made the resident-Q4_K
16// decode 3x SLOWER than the i16 cache (1.5B: 758 vs 254 ms/token) -- exactly the "scalar unpack dominates" verdict
17// nx_nofloat_w4_gemm_gate measured. One u64 read yields 8 bytes; the low and high nibbles are masked in place and
18// each 4-byte half is spread into four i16 lanes with two shift-or-mask steps: 22 ops per 8 bytes where the byte
19// path spent ~160. Bit-exact by construction (the same lanes end up with the same values; nx_nofloat_q4k_gate
20// re-proves the dot against an exact i64 reference on both shift regimes and the 7B ffn width).
21const DS_M_NIB: i64 = 0x0F0F0F0F0F0F0F0F // low nibble of every byte
22const DS_M_LO32: i64 = 0xFFFFFFFF // the 4 bytes a spread turns into 4 lanes
23const DS_M_SPREAD16: i64 = 0x0000FFFF0000FFFF // after (v | v<<16): bytes 0,1 stay, bytes 2,3 land at 32..47
24const DS_M_SPREAD8: i64 = 0x00FF00FF00FF00FF // after (v | v<<8): one byte per 16-bit lane
25const DS_HALF_BITS: i64 = 32
26const DS_NIB_BITS: i64 = 4
27const DS_QS_U64_PER_GROUP: i64 = 4 // 32 qs bytes per sub-block pair = four u64 reads
28
29// four packed bytes (bits 0..31 of v) -> four i16 lanes (b0 | b1<<16 | b2<<32 | b3<<48), the ds_pack4 layout.
30func ds_spread4(v: i64) -> i64 {
31 var y: i64 = v & DS_M_LO32
32 y = (y | (y << 16)) & DS_M_SPREAD16
33 y = (y | (y << 8)) & DS_M_SPREAD8
34 return y
35}
36
37func ds_pack4(a: i64, b: i64, c: i64, d: i64) -> i64 {
38 return (a & 0xFFFF) | ((b & 0xFFFF) << 16) | ((c & 0xFFFF) << 32) | ((d & 0xFFFF) << 48)
39}
40
41func ds_hsum(acc: *i64) -> i64 {
42 var sum: i64 = 0
43 var i: i64 = 0
44 while i < 4 {
45 let v: i64 = acc[i]
46 var lo: i64 = v & 0xFFFFFFFF
47 if lo >= 0x80000000 { lo = lo - 0x100000000 }
48 var hi: i64 = (v >> 32) & 0xFFFFFFFF
49 if hi >= 0x80000000 { hi = hi - 0x100000000 }
50 sum = sum + lo + hi
51 i = i + 1
52 }
53 return sum
54}
55
56// sc_pre[blk*8 + is] = Σ_{l<32} col[blk*256 + is*32 + l] (activation-only, precompute once)
57func nx_q4k_sc_precompute(col_q10: *i64, n_blocks: i64, sc_pre: *i64) -> i64 {
58 var blk: i64 = 0
59 while blk < n_blocks {
60 var is_: i64 = 0
61 while is_ < 8 {
62 var s: i64 = 0
63 var l: i64 = 0
64 let base: i64 = blk * 256 + is_ * 32
65 while l < 32 { s = s + col_q10[base + l]; l = l + 1 }
66 sc_pre[blk * 8 + is_] = s
67 is_ = is_ + 1
68 }
69 blk = blk + 1
70 }
71 return 0
72}
73
74func nx_q4k_dot_simd(buf: *u8, base_off: i64, n_blocks: i64, col_i16: *i64,
75 qpk: *i64, qhi: *i64, acc: *i64, sc_pre: *i64) -> i64 {
76 var dot: i64 = 0
77 var blk: i64 = 0
78 while blk < n_blocks {
79 let sb: i64 = base_off + blk * 144
80 let d_q24: i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb))
81 let dmin_q24: i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb + 2))
82 let scales_off: i64 = sb + 4
83 let qs_off: i64 = sb + 16
84 let col_base: i64 = blk * 256
85 // SCALE FOLDED INTO THE LANES (LM4b, second cut): the 6-bit sub-block scale multiplies the spread nibble
86 // words directly (sc*q <= 63*15 = 945 per 16-bit lane, no carry into the neighbour), so the eight i32 lanes
87 // accumulate sum over ALL EIGHT sub-blocks of (sc*q)*col and ONE horizontal sum serves the super-block
88 // instead of eight. Exact by algebra: sum_sub d*sc_sub*sq_sub = d * sum_sub (sc_sub*sq_sub). Lane bound:
89 // 32 products per lane per super-block, each <= 945*32767 -> < 9.9e8 < 2^31 even at every worst case at once.
90 acc[0] = 0; acc[1] = 0; acc[2] = 0; acc[3] = 0
91 var msum: i64 = 0 // sum over sub-blocks of m_sub * (sum of the sub-block's activation): the dmin term
92 var g: i64 = 0
93 while g < 4 {
94 let is0: i64 = g + g
95 let is1: i64 = is0 + 1
96 var sc0: i64 = 0
97 var m0: i64 = 0
98 var sc1: i64 = 0
99 var m1s: i64 = 0
100 if is0 < 4 {
101 sc0 = nx_le_read_u8(buf, scales_off + is0) & 0x3F
102 m0 = nx_le_read_u8(buf, scales_off + is0 + 4) & 0x3F
103 } else {
104 let k0: i64 = is0 - 4
105 sc0 = ((nx_le_read_u8(buf, scales_off + k0) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k0) & 0x0F)
106 m0 = ((nx_le_read_u8(buf, scales_off + 4 + k0) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k0) >> 4)
107 }
108 if is1 < 4 {
109 sc1 = nx_le_read_u8(buf, scales_off + is1) & 0x3F
110 m1s = nx_le_read_u8(buf, scales_off + is1 + 4) & 0x3F
111 } else {
112 let k1: i64 = is1 - 4
113 sc1 = ((nx_le_read_u8(buf, scales_off + k1) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k1) & 0x0F)
114 m1s = ((nx_le_read_u8(buf, scales_off + 4 + k1) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k1) >> 4)
115 }
116 let grp: i64 = qs_off + g * 32
117
118 // read 32 qs bytes as four u64; split nibbles in place; spread each 4-byte half into i16 lanes; scale the lanes
119 var j: i64 = 0
120 while j < DS_QS_U64_PER_GROUP {
121 let w: i64 = nx_le_read_u64(buf, grp + j * 8)
122 let lo: i64 = w & DS_M_NIB
123 let hi: i64 = (w >> DS_NIB_BITS) & DS_M_NIB
124 qpk[j + j] = ds_spread4(lo) * sc0
125 qpk[j + j + 1] = ds_spread4(lo >> DS_HALF_BITS) * sc0
126 qhi[j + j] = ds_spread4(hi) * sc1
127 qhi[j + j + 1] = ds_spread4(hi >> DS_HALF_BITS) * sc1
128 j = j + 1
129 }
130
131 let cl0: i64 = (col_i16 as i64) + ((col_base + is0 * 32) / 4) * 8
132 __i16x16_madd(acc as *i64, qpk as *i64, cl0 as *i64)
133 __i16x16_madd(acc as *i64, ((qpk as i64) + 32) as *i64, (cl0 + 32) as *i64)
134 let cl1: i64 = (col_i16 as i64) + ((col_base + is1 * 32) / 4) * 8
135 __i16x16_madd(acc as *i64, qhi as *i64, cl1 as *i64)
136 __i16x16_madd(acc as *i64, ((qhi as i64) + 32) as *i64, (cl1 + 32) as *i64)
137
138 msum = msum + m0 * sc_pre[blk * 8 + is0] + m1s * sc_pre[blk * 8 + is1]
139 g = g + 1
140 }
141 dot = dot + d_q24 * ds_hsum(acc) - dmin_q24 * msum
142 blk = blk + 1
143 }
144 return dot
145}
146