nx_q4k_matmul.nx source

1// nx_q4k_matmul.nx -- fused-dequant Q4_K row × column dot product. 2// 3// Phase C foundation brick (conductor cardinal 2026-05-19). 4// Composes nx_dequant_iter.nx with the existing matmul pattern to 5// produce a row × column dot product WITHOUT materializing a full 6// dequantized weight tensor. 7// 8// What the brick does: 9// * Caller holds the Q4_K-quantized weight buffer (e.g. raw bytes 10// from a GGUF file -- composes with nx_gguf_load_lazy) 11// * Caller holds the dense column vector (Q10 i64) 12// * This brick walks the row's Q4_K blocks, dequants each block's 13// 256 values inline, multiplies by the corresponding column 14// entries, and accumulates the sum 15// * NO scratch buffer for dequantized weights; the substrate-side 16// RAM cost is the column vector + per-block iterator (~152 bytes) 17// 18// Honest framing (per no-strawman-perf-comparisons cardinal): 19// - This brick proves the FUSED dequant+dot composition works and 20// produces bit-exact results vs the materialize-then-matmul path. 21// - Per-element throughput on qemu-RV64 is measured by the bench 22// below. This is NOT a claim about CUDA or llama.cpp -- those 23// systems have their own per-element costs that require running 24// them on the same hardware to compare honestly. 25// 26// Output unit: dot product is returned in Q34 (sum of Q24 weight × 27// Q10 column products). The Q4_K weight dequant is now Q24 (EXACT 28// super-scales -- Q10/Q14 underflow/flip sign on real trained-model 29// super-scales). Caller uses nx_q4km_q20_to_q10 to recover Q10 30// (>>24 round-half). Q34 is kept internally so successive matmuls can 31// compose without intermediate normalization rounding. 32// 33// genealogy_id: ggml_q4k_dot_canon + standard_fused_dequant_pattern 34// lineage_id: substrate_q4k_matmul_v1_fused 35 36import "nx_syscalls.nx" 37import "nx_tier.nx" 38import "nx_le.nx" 39import "nx_gguf.nx" 40import "nx_gguf_load.nx" 41import "nx_dequant_iter.nx" 42 43// ===== Sealed-enum: Q4KMatmulVerdict ============================== 44 45const NX_Q4KM_OK: nx_int = 0 46const NX_Q4KM_ERR_BAD_DIM: nx_int = 1 47const NX_Q4KM_N_VERDICTS: nx_int = 2 48 49func nx_q4km_verdict_is_valid(v: nx_int) -> nx_int { 50 if v < 0 { return 0 } 51 if v >= NX_Q4KM_N_VERDICTS { return 0 } 52 return 1 53} 54 55// ===== Public: fused row × col dot product ========================= 56// 57// weight_buf: byte buffer containing Q4_K weight blocks 58// weight_super_off: byte offset of the row's FIRST Q4_K super-block 59// (subsequent blocks are at +144, +288, ... up 60// to weight_super_off + n_blocks*144) 61// n_blocks: number of Q4_K super-blocks in the row 62// (= K_dim / 256; K_dim must be a multiple of 256) 63// col_q10: dense column vector, K_dim Q10 i64 values 64// iter: caller-owned NxQ4KBlockIter (alloc once, reuse 65// across rows/columns -- the matmul caller's job) 66// 67// Returns: dot product in Q34 (Q24 weight × Q10 col). Caller uses 68// nx_q4km_q20_to_q10 to recover Q10. 69 70func nx_q4k_dot_row_col(weight_buf: *u8, weight_super_off: i64, 71 n_blocks: nx_int, 72 col_q10: *i64, 73 iter: *NxQ4KBlockIter) -> i64 { 74 // NO PER-CALL SCRATCH. This function is called m*n times by a GEMM; the 75 // two sys_mmap(8) cells it used to allocate here for nx_q4k_iter_scale_min 76 // were never unmapped, leaking a page-rounded mapping per call (~9,700 per 77 // matmul at n=4864) -- an unclosed resource in a loop is a clock, not a 78 // leak. nx_q4k_iter_scale/nx_q4k_iter_min return the values directly. 79 var dot: i64 = 0 80 81 var blk: nx_int = 0 82 while blk < n_blocks { 83 let blk_off: i64 = weight_super_off + blk * NX_GL_Q4_K_BPB 84 nx_q4k_iter_init(weight_buf, blk_off, iter) 85 let col_base: i64 = blk * 256 86 87 // ggml layout (verified bit-exact): sub-block sb draws from 88 // 32-byte group g = sb/2; output position l (0..32, in order) = 89 // byte qs[g*32 + l], LOW nibble when sb even, HIGH when sb odd. 90 // Activation column aligns to OUTPUT position col_base + sb*32 + l. 91 var sb: nx_int = 0 92 while sb < 8 { 93 let sc_sb: i64 = nx_q4k_iter_scale(iter, sb) 94 let m_sb: i64 = nx_q4k_iter_min(iter, sb) 95 let d1: i64 = iter.d_q10 * sc_sb // Q24 * sc 96 let m1: i64 = iter.dmin_q10 * m_sb // Q24 * m 97 let g: nx_int = sb / 2 98 let is_high: nx_int = sb - g * 2 // 0 low, 1 high 99 let grp_base: i64 = g * 32 100 let col_sb_base: i64 = col_base + sb * 32 101 102 var l: nx_int = 0 103 while l < 32 { 104 let byte_v: i64 = nx_le_read_u8(iter.qs_ptr, grp_base + l) 105 var q4: i64 = 0 106 if is_high == 0 { 107 q4 = byte_v & 0x0F 108 } else { 109 q4 = byte_v >> 4 110 } 111 let v: i64 = d1 * q4 - m1 // Q24 112 dot = dot + v * col_q10[col_sb_base + l] // Q24 * Q10 = Q34 113 l = l + 1 114 } 115 sb = sb + 1 116 } 117 blk = blk + 1 118 } 119 return dot 120} 121 122// ===== Public: convenience -- normalize fused dot -> Q10 ========== 123// 124// The fused dot now accumulates in Q34 (Q24 weight × Q10 column). 125// Dividing by 2^24 recovers a Q10 scalar. Round-half away from zero 126// (add/subtract 2^23 before the shift), matching the i64 truncation 127// semantics used elsewhere in the substrate. Name kept for callers. 128// 129// NOTE: this is now a >>24 normalization (was /1024 for the old Q20 130// path). Any caller feeding a genuinely Q20 value must not use this. 131 132const NX_Q4KM_Q24: i64 = 16777216 // 2^24 (weight super-scale unit) 133const NX_Q4KM_Q23: i64 = 8388608 // 2^23 (round-half bias) 134 135func nx_q4km_q20_to_q10(dot_q34: i64) -> i64 { 136 if dot_q34 >= 0 { 137 return (dot_q34 + NX_Q4KM_Q23) / NX_Q4KM_Q24 138 } 139 return (dot_q34 - NX_Q4KM_Q23) / NX_Q4KM_Q24 140}