nx_q4k_dot_simd2_lib.nx source
↩ module page · 105 lines · 5851 B
1// nx_q4k_dot_simd2_lib.nx -- the fused Q4_K SIMD dot with the COMPILER-EMITTED unpack (2026-09-02, LM4c).
2// Same algebra as nx_q4k_dot_simd (nx_q4k_dot_simd_lib): per super-block the eight i32 lanes accumulate
3// sum over all eight sub-blocks of (sc_sub*q)*col and ONE horizontal sum serves the super-block; the dmin
4// term is a scalar sum of m_sub times the precomputed sub-block activation sums. The ONLY difference is
5// HOW the 32 packed nibble bytes of a sub-block pair become 64 scaled i16 lanes: nx_q4k_dot_simd does it in
6// scalar NishiLang (four u64 reads, eight bit-spreads, eight 64-bit multiplies -- about 1.4 ops per weight,
7// the measured residual between the resident-Q4_K decode and the i16 cache route), this one calls
8// __q4k_unpack32s, one intrinsic that nx_cc emits as 14 AVX2 instructions (vpmovzxbw, vpsrlw, vpand,
9// vpmullw, vmovups). Bit-identical by construction: every lane ends up holding exactly (nibble * sc).
10// REQUIRES a compiler that knows __q4k_unpack32s (nx_types OP_Q4KUNPACK32S); a consumer built with an
11// older compiler fails to COMPILE, never silently falls back -- which is why this lives in its own lib and
12// nx_q4k_dot_simd_lib stays as the compiler-neutral twin until the toolchain is promoted on every host.
13// qbuf: 128 B caller scratch -- lo lanes at qbuf[0..31], hi lanes at qbuf[32..63] (i16), the layout
14// __i16x16_madd consumes directly (two 32-byte halves per group).
15// license_tier: ORIGINAL
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_le.nx"
19import "nx_tensor.nx"
20import "nx_gguf.nx"
21import "nx_gguf_load.nx"
22import "nx_q4k_dot_simd_lib.nx"
23
24const DS2_SUPER_BYTES: i64 = 144
25const DS2_SUPER_VALS: i64 = 256
26const DS2_SCALES_OFF: i64 = 4
27const DS2_QS_OFF: i64 = 16
28const DS2_GROUP_BYTES: i64 = 32
29const DS2_SUB_VALS: i64 = 32
30const DS2_SUBBLOCKS: i64 = 8
31const DS2_SC_HI_SHIFT: i64 = 16 // scpack = sc_lo | sc_hi << 16
32const DS2_QBUF_HI_OFF: i64 = 64 // hi lanes start 64 bytes into qbuf (32 i16)
33const DS2_HALF_BYTES: i64 = 32 // one __i16x16_madd consumes 32 bytes (16 lanes)
34
35// Same contract as nx_q4k_dot_simd(buf, base_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre) with qpk/qhi
36// replaced by ONE 128-byte qbuf. Returns sum over the row of W_Q24 * col_i16.
37func nx_q4k_dot_simd2(buf: *u8, base_off: i64, n_blocks: i64, col_i16: *i64, qbuf: *i64, acc: *i64, sc_pre: *i64) -> i64 {
38 var dot: i64 = 0
39 var blk: i64 = 0
40 let qlo: i64 = qbuf as i64
41 let qhi: i64 = qlo + DS2_QBUF_HI_OFF
42 while blk < n_blocks {
43 let sb: i64 = base_off + blk * DS2_SUPER_BYTES
44 let d_q24: i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb))
45 let dmin_q24: i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb + 2))
46 let scales_off: i64 = sb + DS2_SCALES_OFF
47 let qs_off: i64 = sb + DS2_QS_OFF
48 let col_base: i64 = blk * DS2_SUPER_VALS
49 acc[0] = 0; acc[1] = 0; acc[2] = 0; acc[3] = 0
50 var msum: i64 = 0
51 // SCALE DECODE FROM TWO WIDE READS (LM4c second cut). MEASURED by the serve profile + the pool probe: with the
52 // unpack in the compiler, a super-block still spent most of its cycles in the 24 nx_le_read_u8 CALLS that
53 // decode the packed 6-bit scales and mins. The 12 scale bytes are read once (u64 + u32) and every field is a
54 // shift-and-mask on registers: byte j = (w0 >> 8j) & 255 for j < 8, byte 8+k = (w1 >> 8k) & 255.
55 let w0: i64 = nx_le_read_u64(buf, scales_off)
56 let w1: i64 = nx_le_read_u32(buf, scales_off + 8)
57 var g: i64 = 0
58 while g < 4 {
59 let is0: i64 = g + g
60 let is1: i64 = is0 + 1
61 var sc0: i64 = 0
62 var m0: i64 = 0
63 var sc1: i64 = 0
64 var m1s: i64 = 0
65 if is0 < 4 {
66 // sub-blocks 0..3: sc = b[is] & 63, m = b[is+4] & 63
67 sc0 = (w0 >> (is0 * 8)) & 0x3F
68 m0 = (w0 >> ((is0 + 4) * 8)) & 0x3F
69 sc1 = (w0 >> (is1 * 8)) & 0x3F
70 m1s = (w0 >> ((is1 + 4) * 8)) & 0x3F
71 } else {
72 // sub-blocks 4..7 (k = is - 4): sc = (b[k]>>6)<<4 | (b[8+k] & 15), m = (b[4+k]>>6)<<4 | (b[8+k] >> 4)
73 let k0: i64 = is0 - 4
74 let k1: i64 = is1 - 4
75 let b0: i64 = (w0 >> (k0 * 8)) & 0xFF
76 let b4: i64 = (w0 >> ((k0 + 4) * 8)) & 0xFF
77 let b8: i64 = (w1 >> (k0 * 8)) & 0xFF
78 sc0 = ((b0 >> 6) << 4) | (b8 & 0x0F)
79 m0 = ((b4 >> 6) << 4) | (b8 >> 4)
80 let c0: i64 = (w0 >> (k1 * 8)) & 0xFF
81 let c4: i64 = (w0 >> ((k1 + 4) * 8)) & 0xFF
82 let c8: i64 = (w1 >> (k1 * 8)) & 0xFF
83 sc1 = ((c0 >> 6) << 4) | (c8 & 0x0F)
84 m1s = ((c4 >> 6) << 4) | (c8 >> 4)
85 }
86 let grp: i64 = qs_off + g * DS2_GROUP_BYTES
87
88 // ONE intrinsic: 32 packed bytes -> 64 scaled i16 lanes (lo * sc0 at qbuf[0..31], hi * sc1 at qbuf[32..63])
89 __q4k_unpack32s(((buf as i64) + grp) as *u8, qbuf as *u8, sc0 | (sc1 << DS2_SC_HI_SHIFT))
90
91 let cl0: i64 = (col_i16 as i64) + ((col_base + is0 * DS2_SUB_VALS) / 4) * 8
92 __i16x16_madd(acc as *i64, qlo as *i64, cl0 as *i64)
93 __i16x16_madd(acc as *i64, (qlo + DS2_HALF_BYTES) as *i64, (cl0 + DS2_HALF_BYTES) as *i64)
94 let cl1: i64 = (col_i16 as i64) + ((col_base + is1 * DS2_SUB_VALS) / 4) * 8
95 __i16x16_madd(acc as *i64, qhi as *i64, cl1 as *i64)
96 __i16x16_madd(acc as *i64, (qhi + DS2_HALF_BYTES) as *i64, (cl1 + DS2_HALF_BYTES) as *i64)
97
98 msum = msum + m0 * sc_pre[blk * DS2_SUBBLOCKS + is0] + m1s * sc_pre[blk * DS2_SUBBLOCKS + is1]
99 g = g + 1
100 }
101 dot = dot + d_q24 * ds_hsum(acc) - dmin_q24 * msum
102 blk = blk + 1
103 }
104 return dot
105}